refactor(web): asset grid state (#3513)

* refactor(web): asset grid state

* fix: multi-select across time buckets

---------

Co-authored-by: Alex Tran <alex.tran1502@gmail.com>

web/src/lib/components/photos-page/actions/select-all-assets.svelte

@@ -19,7 +19,7 @@
 
       const assetGridState = get(assetStore);
       for (const bucket of assetGridState.buckets) {
-        await assetStore.getAssetsByBucket(bucket.bucketDate, BucketPosition.Unknown);
+        await assetStore.loadBucket(bucket.bucketDate, BucketPosition.Unknown);
         for (const asset of bucket.assets) {
           assetInteractionStore.addAssetToMultiselectGroup(asset);
         }

web/src/lib/components/photos-page/asset-date-group.svelte

@@ -60,7 +60,7 @@
 
   $: {
     if (actualBucketHeight && actualBucketHeight != 0 && actualBucketHeight != bucketHeight) {
-      const heightDelta = assetStore.updateBucketHeight(bucketDate, actualBucketHeight);
+      const heightDelta = assetStore.updateBucket(bucketDate, actualBucketHeight);
       if (heightDelta !== 0) {
         scrollTimeline(heightDelta);
       }

web/src/lib/components/photos-page/asset-grid.svelte

@@ -4,7 +4,7 @@
   import { locale } from '$lib/stores/preferences.store';
   import { formatGroupTitle, splitBucketIntoDateGroups } from '$lib/utils/timeline-util';
   import type { UserResponseDto } from '@api';
-  import { api, AssetCountByTimeBucketResponseDto, AssetResponseDto, TimeGroupEnum } from '@api';
+  import { AssetResponseDto, TimeGroupEnum, api } from '@api';
   import { DateTime } from 'luxon';
   import { onDestroy, onMount } from 'svelte';
   import AssetViewer from '../asset-viewer/asset-viewer.svelte';
@@ -17,13 +17,13 @@
   import AssetDateGroup from './asset-date-group.svelte';
   import MemoryLane from './memory-lane.svelte';
 
-  import { AppRoute } from '$lib/constants';
-  import { goto } from '$app/navigation';
   import { browser } from '$app/environment';
+  import { goto } from '$app/navigation';
+  import { AppRoute } from '$lib/constants';
+  import type { AssetInteractionStore } from '$lib/stores/asset-interaction.store';
+  import type { AssetStore } from '$lib/stores/assets.store';
   import { isSearchEnabled } from '$lib/stores/search.store';
   import ShowShortcuts from '../shared-components/show-shortcuts.svelte';
-  import type { AssetStore } from '$lib/stores/assets.store';
-  import type { AssetInteractionStore } from '$lib/stores/asset-interaction.store';
 
   export let user: UserResponseDto | undefined = undefined;
   export let isAlbumSelectionMode = false;
@@ -39,14 +39,13 @@
   let viewportHeight = 0;
   let viewportWidth = 0;
   let assetGridElement: HTMLElement;
-  let bucketInfo: AssetCountByTimeBucketResponseDto;
   let showShortcuts = false;
 
   const onKeyboardPress = (event: KeyboardEvent) => handleKeyboardPress(event);
 
   onMount(async () => {
     document.addEventListener('keydown', onKeyboardPress);
-    const { data: assetCountByTimebucket } = await api.assetApi.getAssetCountByTimeBucket({
+    const { data: timeBuckets } = await api.assetApi.getAssetCountByTimeBucket({
       getAssetCountByTimeBucketDto: {
         timeGroup: TimeGroupEnum.Month,
         userId: user?.id,
@@ -54,26 +53,7 @@
       },
     });
 
-    bucketInfo = assetCountByTimebucket;
-
-    assetStore.setInitialState(viewportHeight, viewportWidth, assetCountByTimebucket, user?.id);
-
-    // Get asset bucket if bucket height is smaller than viewport height
-    let bucketsToFetchInitially: string[] = [];
-    let initialBucketsHeight = 0;
-    $assetStore.buckets.every((bucket) => {
-      if (initialBucketsHeight < viewportHeight) {
-        initialBucketsHeight += bucket.bucketHeight;
-        bucketsToFetchInitially.push(bucket.bucketDate);
-        return true;
-      } else {
-        return false;
-      }
-    });
-
-    bucketsToFetchInitially.forEach((bucketDate) => {
-      assetStore.getAssetsByBucket(bucketDate, BucketPosition.Visible);
-    });
+    assetStore.init({ width: viewportHeight, height: viewportWidth }, timeBuckets.buckets, user?.id);
   });
 
   onDestroy(() => {
@@ -81,7 +61,7 @@
       document.removeEventListener('keydown', onKeyboardPress);
     }
 
-    assetStore.setInitialState(0, 0, { totalCount: 0, buckets: [] }, undefined);
+    assetStore.init({ width: 0, height: 0 }, [], undefined);
   });
 
   const handleKeyboardPress = (event: KeyboardEvent) => {
@@ -113,7 +93,7 @@
     const target = el.firstChild as HTMLElement;
     if (target) {
       const bucketDate = target.id.split('_')[1];
-      assetStore.getAssetsByBucket(bucketDate, event.detail.position);
+      assetStore.loadBucket(bucketDate, event.detail.position);
     }
   }
 
@@ -122,14 +102,14 @@
   }
 
   const navigateToPreviousAsset = async () => {
-    const prevAsset = await assetStore.getAdjacentAsset($viewingAsset.id, 'previous');
+    const prevAsset = await assetStore.getPreviousAssetId($viewingAsset.id);
     if (prevAsset) {
       assetViewingStore.setAssetId(prevAsset);
     }
   };
 
   const navigateToNextAsset = async () => {
-    const nextAsset = await assetStore.getAdjacentAsset($viewingAsset.id, 'next');
+    const nextAsset = await assetStore.getNextAssetId($viewingAsset.id);
     if (nextAsset) {
       assetViewingStore.setAssetId(nextAsset);
     }
@@ -234,8 +214,12 @@
     assetInteractionStore.clearAssetSelectionCandidates();
 
     if ($assetSelectionStart && rangeSelection) {
-      let startBucketIndex = $assetStore.loadedAssets[$assetSelectionStart.id];
-      let endBucketIndex = $assetStore.loadedAssets[asset.id];
+      let startBucketIndex = $assetStore.getBucketIndexByAssetId($assetSelectionStart.id);
+      let endBucketIndex = $assetStore.getBucketIndexByAssetId(asset.id);
+
+      if (startBucketIndex === null || endBucketIndex === null) {
+        return;
+      }
 
       if (endBucketIndex < startBucketIndex) {
         [startBucketIndex, endBucketIndex] = [endBucketIndex, startBucketIndex];
@@ -244,7 +228,7 @@
       // Select/deselect assets in all intermediate buckets
       for (let bucketIndex = startBucketIndex + 1; bucketIndex < endBucketIndex; bucketIndex++) {
         const bucket = $assetStore.buckets[bucketIndex];
-        await assetStore.getAssetsByBucket(bucket.bucketDate, BucketPosition.Unknown);
+        await assetStore.loadBucket(bucket.bucketDate, BucketPosition.Unknown);
         for (const asset of bucket.assets) {
           if (deselect) {
             assetInteractionStore.removeAssetFromMultiselectGroup(asset);
@@ -308,7 +292,7 @@
   <ShowShortcuts on:close={() => (showShortcuts = !showShortcuts)} />
 {/if}
 
-{#if bucketInfo && viewportHeight && $assetStore.timelineHeight > viewportHeight}
+{#if viewportHeight && $assetStore.initialized && $assetStore.timelineHeight > viewportHeight}
   <Scrollbar
     {assetStore}
     scrollbarHeight={viewportHeight}
@@ -335,9 +319,7 @@
       {#each $assetStore.buckets as bucket, bucketIndex (bucketIndex)}
         <IntersectionObserver
           on:intersected={intersectedHandler}
-          on:hidden={async () => {
-            await assetStore.cancelBucketRequest(bucket.cancelToken, bucket.bucketDate);
-          }}
+          on:hidden={() => assetStore.cancelBucket(bucket)}
           let:intersecting
           top={750}
           bottom={750}

web/src/lib/models/asset-grid-state.ts

@@ -1,4 +1,7 @@
-import type { AssetResponseDto } from '@api';
+import { api, AssetCountByTimeBucket, AssetResponseDto } from '@api';
+import { writable } from 'svelte/store';
+import type { AssetStore } from '../stores/assets.store';
+import { handleError } from '../utils/handle-error';
 
 export enum BucketPosition {
   Above = 'above',
@@ -7,6 +10,17 @@ export enum BucketPosition {
   Unknown = 'unknown',
 }
 
+export interface Viewport {
+  width: number;
+  height: number;
+}
+
+interface AssetLookup {
+  bucket: AssetBucket;
+  bucketIndex: number;
+  assetIndex: number;
+}
+
 export class AssetBucket {
   /**
    * The DOM height of the bucket in pixel
@@ -15,44 +29,220 @@ export class AssetBucket {
   bucketHeight!: number;
   bucketDate!: string;
   assets!: AssetResponseDto[];
-  cancelToken!: AbortController;
+  cancelToken!: AbortController | null;
   position!: BucketPosition;
 }
 
-export class AssetGridState {
-  /**
-   * The total height of the timeline in pixel
-   * This value is first estimated by the number of asset and later is corrected as the user scroll
-   */
+const THUMBNAIL_HEIGHT = 235;
+
+export class AssetGridState implements AssetStore {
+  private store$ = writable(this);
+  private assetToBucket: Record<string, AssetLookup> = {};
+  private viewport: Viewport = { width: 0, height: 0 };
+  private userId: string | undefined;
+
+  initialized = false;
   timelineHeight = 0;
+  buckets: AssetBucket[] = [];
+  assets: AssetResponseDto[] = [];
 
-  /**
-   * The fixed viewport height in pixel
-   */
-  viewportHeight = 0;
+  subscribe = this.store$.subscribe;
 
-  /**
-   * The fixed viewport width in pixel
-   */
-  viewportWidth = 0;
+  init(viewport: Viewport, buckets: AssetCountByTimeBucket[], userId: string | undefined) {
+    this.initialized = false;
+    this.assets = [];
+    this.assetToBucket = {};
+    this.buckets = [];
+    this.viewport = viewport;
+    this.userId = userId;
+    this.buckets = buckets.map((bucket) => {
+      const unwrappedWidth = (3 / 2) * bucket.count * THUMBNAIL_HEIGHT * (7 / 10);
+      const rows = Math.ceil(unwrappedWidth / this.viewport.width);
+      const height = rows * THUMBNAIL_HEIGHT;
 
-  /**
-   * List of bucket information
-   */
-  buckets: AssetBucket[] = [];
+      return {
+        bucketDate: bucket.timeBucket,
+        bucketHeight: height,
+        assets: [],
+        cancelToken: null,
+        position: BucketPosition.Unknown,
+      };
+    });
 
-  /**
-   * Total assets that have been loaded
-   */
-  assets: AssetResponseDto[] = [];
+    this.timelineHeight = this.buckets.reduce((acc, b) => acc + b.bucketHeight, 0);
 
-  /**
-   * Total assets that have been loaded along with additional data
-   */
-  loadedAssets: Record<string, number> = {};
+    this.emit(false);
 
-  /**
-   * User that owns assets
-   */
-  userId: string | undefined;
+    let height = 0;
+    for (const bucket of this.buckets) {
+      if (height < this.viewport.height) {
+        height += bucket.bucketHeight;
+        this.loadBucket(bucket.bucketDate, BucketPosition.Visible);
+        continue;
+      }
+
+      break;
+    }
+
+    this.initialized = true;
+  }
+
+  getBucketByDate(bucketDate: string): AssetBucket | null {
+    return this.buckets.find((bucket) => bucket.bucketDate === bucketDate) || null;
+  }
+
+  getBucketInfoForAssetId(assetId: string) {
+    return this.assetToBucket[assetId] || null;
+  }
+
+  getBucketIndexByAssetId(assetId: string) {
+    return this.assetToBucket[assetId]?.bucketIndex ?? null;
+  }
+
+  async loadBucket(bucketDate: string, position: BucketPosition): Promise<void> {
+    try {
+      const bucket = this.getBucketByDate(bucketDate);
+      if (!bucket) {
+        return;
+      }
+
+      bucket.position = position;
+
+      if (bucket.assets.length !== 0) {
+        this.emit(false);
+        return;
+      }
+
+      bucket.cancelToken = new AbortController();
+
+      const { data: assets } = await api.assetApi.getAssetByTimeBucket(
+        {
+          getAssetByTimeBucketDto: {
+            timeBucket: [bucketDate],
+            userId: this.userId,
+            withoutThumbs: true,
+          },
+        },
+        { signal: bucket.cancelToken.signal },
+      );
+
+      bucket.assets = assets;
+      this.emit(true);
+    } catch (error) {
+      handleError(error, 'Failed to load assets');
+    }
+  }
+
+  cancelBucket(bucket: AssetBucket) {
+    bucket.cancelToken?.abort();
+  }
+
+  updateBucket(bucketDate: string, height: number) {
+    const bucket = this.getBucketByDate(bucketDate);
+    if (!bucket) {
+      return 0;
+    }
+
+    const delta = height - bucket.bucketHeight;
+    const scrollTimeline = bucket.position == BucketPosition.Above;
+
+    bucket.bucketHeight = height;
+    bucket.position = BucketPosition.Unknown;
+
+    this.timelineHeight += delta;
+
+    this.emit(false);
+
+    return scrollTimeline ? delta : 0;
+  }
+
+  updateAsset(assetId: string, isFavorite: boolean) {
+    const asset = this.assets.find((asset) => asset.id === assetId);
+    if (!asset) {
+      return;
+    }
+
+    asset.isFavorite = isFavorite;
+    this.emit(false);
+  }
+
+  removeAsset(assetId: string) {
+    for (let i = 0; i < this.buckets.length; i++) {
+      const bucket = this.buckets[i];
+      for (let j = 0; j < bucket.assets.length; j++) {
+        const asset = bucket.assets[j];
+        if (asset.id !== assetId) {
+          continue;
+        }
+
+        bucket.assets.splice(j, 1);
+        if (bucket.assets.length === 0) {
+          this.buckets.splice(i, 1);
+        }
+
+        this.emit(true);
+        return;
+      }
+    }
+  }
+
+  async getPreviousAssetId(assetId: string): Promise<string | null> {
+    const info = this.getBucketInfoForAssetId(assetId);
+    if (!info) {
+      return null;
+    }
+
+    const { bucket, assetIndex, bucketIndex } = info;
+
+    if (assetIndex !== 0) {
+      return bucket.assets[assetIndex - 1].id;
+    }
+
+    if (bucketIndex === 0) {
+      return null;
+    }
+
+    const previousBucket = this.buckets[bucketIndex - 1];
+    await this.loadBucket(previousBucket.bucketDate, BucketPosition.Unknown);
+    return previousBucket.assets.at(-1)?.id || null;
+  }
+
+  async getNextAssetId(assetId: string): Promise<string | null> {
+    const info = this.getBucketInfoForAssetId(assetId);
+    if (!info) {
+      return null;
+    }
+
+    const { bucket, assetIndex, bucketIndex } = info;
+
+    if (assetIndex !== bucket.assets.length - 1) {
+      return bucket.assets[assetIndex + 1].id;
+    }
+
+    if (bucketIndex === this.buckets.length - 1) {
+      return null;
+    }
+
+    const nextBucket = this.buckets[bucketIndex + 1];
+    await this.loadBucket(nextBucket.bucketDate, BucketPosition.Unknown);
+    return nextBucket.assets[0]?.id || null;
+  }
+
+  private emit(recalculate: boolean) {
+    if (recalculate) {
+      this.assets = this.buckets.flatMap(({ assets }) => assets);
+
+      const assetToBucket: Record<string, AssetLookup> = {};
+      for (let i = 0; i < this.buckets.length; i++) {
+        const bucket = this.buckets[i];
+        for (let j = 0; j < bucket.assets.length; j++) {
+          const asset = bucket.assets[j];
+          assetToBucket[asset.id] = { bucket, bucketIndex: i, assetIndex: j };
+        }
+      }
+      this.assetToBucket = assetToBucket;
+    }
+
+    this.store$.update(() => this);
+  }
 }

web/src/lib/stores/assets.store.ts

@@ -1,268 +1,38 @@
-import { AssetGridState, BucketPosition } from '$lib/models/asset-grid-state';
-import { api, AssetCountByTimeBucketResponseDto, AssetResponseDto } from '@api';
-import { writable } from 'svelte/store';
+import { AssetBucket, AssetGridState, BucketPosition, Viewport } from '$lib/models/asset-grid-state';
+import type { AssetCountByTimeBucket } from '@api';
 
 export interface AssetStore {
-  setInitialState: (
-    viewportHeight: number,
-    viewportWidth: number,
-    data: AssetCountByTimeBucketResponseDto,
-    userId: string | undefined,
-  ) => void;
-  getAssetsByBucket: (bucket: string, position: BucketPosition) => Promise<void>;
-  updateBucketHeight: (bucket: string, actualBucketHeight: number) => number;
-  cancelBucketRequest: (token: AbortController, bucketDate: string) => Promise<void>;
-  getAdjacentAsset: (assetId: string, direction: 'next' | 'previous') => Promise<string | null>;
+  init: (viewport: Viewport, data: AssetCountByTimeBucket[], userId: string | undefined) => void;
+
+  // bucket
+  loadBucket: (bucket: string, position: BucketPosition) => Promise<void>;
+  updateBucket: (bucket: string, actualBucketHeight: number) => number;
+  cancelBucket: (bucket: AssetBucket) => void;
+
+  // asset
   removeAsset: (assetId: string) => void;
   updateAsset: (assetId: string, isFavorite: boolean) => void;
+
+  // asset navigation
+  getNextAssetId: (assetId: string) => Promise<string | null>;
+  getPreviousAssetId: (assetId: string) => Promise<string | null>;
+
+  // store
   subscribe: (run: (value: AssetGridState) => void, invalidate?: (value?: AssetGridState) => void) => () => void;
 }
 
 export function createAssetStore(): AssetStore {
-  let _loadingBuckets: { [key: string]: boolean } = {};
-  let _assetGridState = new AssetGridState();
-
-  const { subscribe, set, update } = writable(new AssetGridState());
-
-  subscribe((state) => {
-    _assetGridState = state;
-  });
-
-  const _estimateViewportHeight = (assetCount: number, viewportWidth: number): number => {
-    // Ideally we would use the average aspect ratio for the photoset, however assume
-    // a normal landscape aspect ratio of 3:2, then discount for the likelihood we
-    // will be scaling down and coalescing.
-    const thumbnailHeight = 235;
-    const unwrappedWidth = (3 / 2) * assetCount * thumbnailHeight * (7 / 10);
-    const rows = Math.ceil(unwrappedWidth / viewportWidth);
-    const height = rows * thumbnailHeight;
-    return height;
-  };
-
-  const refreshLoadedAssets = (state: AssetGridState): void => {
-    state.loadedAssets = {};
-    state.buckets.forEach((bucket, bucketIndex) =>
-      bucket.assets.map((asset) => {
-        state.loadedAssets[asset.id] = bucketIndex;
-      }),
-    );
-  };
-
-  const setInitialState = (
-    viewportHeight: number,
-    viewportWidth: number,
-    data: AssetCountByTimeBucketResponseDto,
-    userId: string | undefined,
-  ) => {
-    set({
-      viewportHeight,
-      viewportWidth,
-      timelineHeight: 0,
-      buckets: data.buckets.map((bucket) => ({
-        bucketDate: bucket.timeBucket,
-        bucketHeight: _estimateViewportHeight(bucket.count, viewportWidth),
-        assets: [],
-        cancelToken: new AbortController(),
-        position: BucketPosition.Unknown,
-      })),
-      assets: [],
-      loadedAssets: {},
-      userId,
-    });
-
-    update((state) => {
-      state.timelineHeight = state.buckets.reduce((acc, b) => acc + b.bucketHeight, 0);
-      return state;
-    });
-  };
-
-  const getAssetsByBucket = async (bucket: string, position: BucketPosition) => {
-    try {
-      const currentBucketData = _assetGridState.buckets.find((b) => b.bucketDate === bucket);
-      if (currentBucketData?.assets && currentBucketData.assets.length > 0) {
-        update((state) => {
-          const bucketIndex = state.buckets.findIndex((b) => b.bucketDate === bucket);
-          state.buckets[bucketIndex].position = position;
-          return state;
-        });
-        return;
-      }
-
-      _loadingBuckets = { ..._loadingBuckets, [bucket]: true };
-      const { data: assets } = await api.assetApi.getAssetByTimeBucket(
-        {
-          getAssetByTimeBucketDto: {
-            timeBucket: [bucket],
-            userId: _assetGridState.userId,
-            withoutThumbs: true,
-          },
-        },
-        { signal: currentBucketData?.cancelToken.signal },
-      );
-      _loadingBuckets = { ..._loadingBuckets, [bucket]: false };
-
-      update((state) => {
-        const bucketIndex = state.buckets.findIndex((b) => b.bucketDate === bucket);
-        state.buckets[bucketIndex].assets = assets;
-        state.buckets[bucketIndex].position = position;
-        state.assets = state.buckets.flatMap((b) => b.assets);
-        refreshLoadedAssets(state);
-        return state;
-      });
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    } catch (e: any) {
-      if (e.name === 'CanceledError') {
-        return;
-      }
-      console.error('Failed to get asset for bucket ', bucket);
-      console.error(e);
-    }
-  };
-
-  const removeAsset = (assetId: string) => {
-    update((state) => {
-      const bucketIndex = state.buckets.findIndex((b) => b.assets.some((a) => a.id === assetId));
-      const assetIndex = state.buckets[bucketIndex].assets.findIndex((a) => a.id === assetId);
-      state.buckets[bucketIndex].assets.splice(assetIndex, 1);
-
-      if (state.buckets[bucketIndex].assets.length === 0) {
-        _removeBucket(state.buckets[bucketIndex].bucketDate);
-      }
-      state.assets = state.buckets.flatMap((b) => b.assets);
-      refreshLoadedAssets(state);
-      return state;
-    });
-  };
-
-  const _removeBucket = (bucketDate: string) => {
-    update((state) => {
-      const bucketIndex = state.buckets.findIndex((b) => b.bucketDate === bucketDate);
-      state.buckets.splice(bucketIndex, 1);
-      state.assets = state.buckets.flatMap((b) => b.assets);
-      refreshLoadedAssets(state);
-      return state;
-    });
-  };
-
-  const updateBucketHeight = (bucket: string, actualBucketHeight: number): number => {
-    let scrollTimeline = false;
-    let heightDelta = 0;
-
-    update((state) => {
-      const bucketIndex = state.buckets.findIndex((b) => b.bucketDate === bucket);
-      // Update timeline height based on the new bucket height
-      const estimateBucketHeight = state.buckets[bucketIndex].bucketHeight;
-
-      heightDelta = actualBucketHeight - estimateBucketHeight;
-      state.timelineHeight += heightDelta;
-
-      scrollTimeline = state.buckets[bucketIndex].position == BucketPosition.Above;
-
-      state.buckets[bucketIndex].bucketHeight = actualBucketHeight;
-      state.buckets[bucketIndex].position = BucketPosition.Unknown;
-
-      return state;
-    });
-
-    if (scrollTimeline) {
-      return heightDelta;
-    }
-
-    return 0;
-  };
-
-  const cancelBucketRequest = async (token: AbortController, bucketDate: string) => {
-    if (!_loadingBuckets[bucketDate]) {
-      return;
-    }
-
-    token.abort();
-
-    update((state) => {
-      const bucketIndex = state.buckets.findIndex((b) => b.bucketDate === bucketDate);
-      state.buckets[bucketIndex].cancelToken = new AbortController();
-      return state;
-    });
-  };
-
-  const updateAsset = (assetId: string, isFavorite: boolean) => {
-    update((state) => {
-      const bucketIndex = state.buckets.findIndex((b) => b.assets.some((a) => a.id === assetId));
-      const assetIndex = state.buckets[bucketIndex].assets.findIndex((a) => a.id === assetId);
-      state.buckets[bucketIndex].assets[assetIndex].isFavorite = isFavorite;
-
-      state.assets = state.buckets.flatMap((b) => b.assets);
-      refreshLoadedAssets(state);
-      return state;
-    });
-  };
-
-  const _getNextAsset = async (currentBucketIndex: number, assetId: string): Promise<AssetResponseDto | null> => {
-    const currentBucket = _assetGridState.buckets[currentBucketIndex];
-    const assetIndex = currentBucket.assets.findIndex(({ id }) => id == assetId);
-    if (assetIndex === -1) {
-      return null;
-    }
-
-    if (assetIndex + 1 < currentBucket.assets.length) {
-      return currentBucket.assets[assetIndex + 1];
-    }
-
-    const nextBucketIndex = currentBucketIndex + 1;
-    if (nextBucketIndex >= _assetGridState.buckets.length) {
-      return null;
-    }
-
-    const nextBucket = _assetGridState.buckets[nextBucketIndex];
-    await getAssetsByBucket(nextBucket.bucketDate, BucketPosition.Unknown);
-
-    return nextBucket.assets[0] ?? null;
-  };
-
-  const _getPrevAsset = async (currentBucketIndex: number, assetId: string): Promise<AssetResponseDto | null> => {
-    const currentBucket = _assetGridState.buckets[currentBucketIndex];
-    const assetIndex = currentBucket.assets.findIndex(({ id }) => id == assetId);
-    if (assetIndex === -1) {
-      return null;
-    }
-
-    if (assetIndex > 0) {
-      return currentBucket.assets[assetIndex - 1];
-    }
-
-    const prevBucketIndex = currentBucketIndex - 1;
-    if (prevBucketIndex < 0) {
-      return null;
-    }
-
-    const prevBucket = _assetGridState.buckets[prevBucketIndex];
-    await getAssetsByBucket(prevBucket.bucketDate, BucketPosition.Unknown);
-
-    return prevBucket.assets[prevBucket.assets.length - 1] ?? null;
-  };
-
-  const getAdjacentAsset = async (assetId: string, direction: 'next' | 'previous'): Promise<string | null> => {
-    const currentBucketIndex = _assetGridState.loadedAssets[assetId];
-    if (currentBucketIndex < 0 || currentBucketIndex >= _assetGridState.buckets.length) {
-      return null;
-    }
-
-    const asset =
-      direction === 'next'
-        ? await _getNextAsset(currentBucketIndex, assetId)
-        : await _getPrevAsset(currentBucketIndex, assetId);
-
-    return asset?.id ?? null;
-  };
+  const store = new AssetGridState();
 
   return {
-    setInitialState,
-    getAssetsByBucket,
-    removeAsset,
-    updateBucketHeight,
-    cancelBucketRequest,
-    getAdjacentAsset,
-    updateAsset,
-    subscribe,
+    init: store.init.bind(store),
+    loadBucket: store.loadBucket.bind(store),
+    updateBucket: store.updateBucket.bind(store),
+    cancelBucket: store.cancelBucket.bind(store),
+    removeAsset: store.removeAsset.bind(store),
+    updateAsset: store.updateAsset.bind(store),
+    getNextAssetId: store.getNextAssetId.bind(store),
+    getPreviousAssetId: store.getPreviousAssetId.bind(store),
+    subscribe: store.subscribe,
   };
 }