[SKStore] KeySets

skiplang/prelude/src/skstore/Context.sk

@@ -341,20 +341,18 @@ mutable class Context private {
   mutable newDirs: SortedSet<DirName> = SortedSet[],
   mutable globals: SortedMap<String, File> = SortedMap[],
   private mutable persistents: SortedMap<String, File> = SortedMap[],
-  private mutable dirty: SortedMap<DirName, SortedSet<Key>> = SortedMap[],
+  private mutable dirty: SortedMap<DirName, KeySet> = SortedMap[],
   private mutable dirtyReaders: /* parentName => childName => keys */ SortedMap<
     DirName,
-    SortedMap<DirName, SortedSet<Key>>,
+    SortedMap<DirName, KeySet>,
   > = SortedMap[],
   mutable canReuse: CanReuse = CRIfMatch(),
   private mutable arrowStack: List<(ArrowKey, TimeStack)> = List[],
-  private mutable lazyGets: SortedMap<DirName, SortedSet<Key>> = SortedMap[],
-  private mutable lazyGetsQueue: Queue<
-    SortedMap<DirName, SortedSet<Key>>,
-  > = Queue[],
+  private mutable lazyGets: SortedMap<DirName, KeySet> = SortedMap[],
+  private mutable lazyGetsQueue: Queue<SortedMap<DirName, KeySet>> = Queue[],
   private mutable lazyGetsRefCount: SortedMap<
     DirName,
-    SortedMap<Key, Int>,
+    KeyMap<Int>,
   > = SortedMap[],
   private mutable sessions: SortedMap<Int, Sub> = SortedMap[],
   private mutable postponables: List<Postponable> = List[],
@@ -593,7 +591,10 @@ mutable class Context private {
         );
         map.set(
           reader.childName,
-          map.maybeGet(reader.childName).default(SortedSet[]).set(reader.key),
+          parent.addReaderKeyToKeySetOpt(
+            reader,
+            map.maybeGet(reader.childName),
+          ),
         )
       };
       time = if (reader.parentName == reader.childName) {
@@ -649,19 +650,20 @@ mutable class Context private {
     this.!lazyGets = SortedMap[]
   }
 
-  mutable fun addLazyGet(dir: LazyDir, key: Key): void {
+  mutable fun addLazyGet(dir: TLazyDir, key: Key): void {
     dirName = dir.dirName;
-    this.!lazyGets[dirName] = this.lazyGets.maybeGet(dirName)
-      .default(SortedSet[])
-      .set(key)
+    this.!lazyGets[dirName] = dir.addKeyToKeySetOpt(
+      key,
+      this.lazyGets.maybeGet(dirName),
+    )
   }
 
-  mutable fun addDirty(dir: EagerDir, key: Key): void {
+  mutable fun addDirty(dir: TEagerDir, key: Key): void {
     dirName = dir.dirName;
-    this.!dirty[dirName] = this.dirty.maybeGet(dirName) match {
-    | None() -> SortedSet[key]
-    | Some(set) -> set.set(key)
-    }
+    this.!dirty[dirName] = dir.addKeyToKeySetOpt(
+      key,
+      this.dirty.maybeGet(dirName),
+    )
   }
 
   mutable fun update(): void {

skiplang/prelude/src/skstore/Dir.sk

@@ -27,6 +27,8 @@ base class Dir protected {timeStack: TimeStack, dirName: DirName} {
   }
 
   fun typed(): TDir;
+
+  fun addReaderKeyToKeySetOpt(arrowKey: ArrowKey, keySetOpt: ?KeySet): KeySet;
 }
 
 base class TDir extends Dir {
@@ -54,6 +56,30 @@ base class TDir extends Dir {
    * Returns the list of files in a directory. Should only be used for testing.
    */
   fun keys(): SortedSet<Key>;
+
+  fun castKeySet(keySet: KeySet): TKeySet<Key> {
+    Unsafe.cast(keySet, TKeySet<Key>)
+  }
+
+  fun castOrEmptyKeySet(keySetOpt: ?KeySet): TKeySet<Key> {
+    keySetOpt.map(this.castKeySet).default(TKeySet::empty())
+  }
+
+  fun addKeyToKeySetOpt(key: Key, keySetOpt: ?KeySet): KeySet {
+    this.castOrEmptyKeySet(keySetOpt).set(key)
+  }
+
+  fun addReaderKeyToKeySetOpt(arrowKey: ArrowKey, keySetOpt: ?KeySet): KeySet {
+    this.addKeyToKeySetOpt(arrowKey.key, keySetOpt)
+  }
+
+  fun castKeyMap<V: frozen>(keyMap: KeyMap<V>): TKeyMap<Key, V> {
+    Unsafe.cast(keyMap, TKeyMap<Key, V>)
+  }
+
+  fun castOrEmptyKeyMap<V: frozen>(keyMapOpt: ?KeyMap<V>): TKeyMap<Key, V> {
+    keyMapOpt.map(this.castKeyMap).default(TKeyMap::empty())
+  }
 }
 
 /*****************************************************************************/

skiplang/prelude/src/skstore/EagerDir.sk

@@ -1420,8 +1420,8 @@ base class EagerDir protected {
   fun update(
     /* this is the parent */
     context: mutable Context,
-    parentDirtyOpt: ?SortedSet<Key>,
-    contextDirtyReadersOpt: ?SortedMap<DirName, SortedSet<Key>>,
+    parentDirtyOpt: ?KeySet,
+    contextDirtyReadersOpt: ?SortedMap<DirName, KeySet>,
     parentMaps: Array<
       (
         MapFun<Key, File>,
@@ -1443,7 +1443,7 @@ base class EagerDir protected {
       contextDirtyReaders.maybeGet(childName) match {
       | None() -> void
       | Some(dirtyKeys) ->
-        keys = dirtyKeys.values();
+        keys = parent.typed().castKeySet(dirtyKeys).values();
         !dirty = withRegionFold(None(), keys, dirty, (_, key, dirtyInRegion) ~> {
           for (p in parentMaps) {
             (_, rangeOpt, _) = p;
@@ -1464,7 +1464,7 @@ base class EagerDir protected {
     parentDirtyOpt match {
     | None() -> void
     | Some(dirtyKeys) ->
-      keys = dirtyKeys.values();
+      keys = parent.typed().castKeySet(dirtyKeys).values();
       !dirty = withRegionFold(None(), keys, dirty, (_, key, dirtyInRegion) ~> {
         for (p in parentMaps) {
           (_, rangeOpt, _) = p;
@@ -1655,7 +1655,7 @@ base class EagerDir protected {
       Some(cdata)
     };
 
-    context.addDirty(this, k);
+    context.addDirty(this.typed(), k);
     path = Path::create(this.dirName, k);
 
     if (this.totalSize != totalSize) {

skiplang/prelude/src/skstore/LazyDir.sk

@@ -21,14 +21,14 @@ base class LazyDir protected {protected collect: Bool} extends Dir {
   /* Updates the lazyGets ref counts and return as well a potentially updated
     LazyDir in which dead keys have been removed. */
   fun updateLazyGets(
-    refCountsOpt: ?SortedMap<Key, Int>,
-    newKeysOpt: ?SortedSet<Key>,
-    deadKeysOpt: ?SortedSet<Key>,
-  ): (?SortedMap<Key, Int>, ?this);
+    refCountsOpt: ?KeyMap<Int>,
+    newKeysOpt: ?KeySet,
+    deadKeysOpt: ?KeySet,
+  ): (?KeyMap<Int>, ?this);
 
   fun update(
     context: mutable Context,
-    dirtyReadersOpt: ?SortedMap<DirName, SortedSet<Key>>,
+    dirtyReadersOpt: ?SortedMap<DirName, KeySet>,
   ): void;
 
   fun typed(): TLazyDir;
@@ -39,18 +39,18 @@ class TLazyDir{
   protected lazyFun: (mutable Context, DirName, Key) ~> ?Array<File>,
 } extends LazyDir, TDir {
   fun updateLazyGets(
-    refCountsOpt: ?SortedMap<Key, Int>,
-    newKeysOpt: ?SortedSet<Key>,
-    deadKeysOpt: ?SortedSet<Key>,
-  ): (?SortedMap<Key, Int>, ?this) {
-    refCounts = refCountsOpt.default(SortedMap[]);
-    newKeys = newKeysOpt.default(SortedSet[]);
-    deadKeys = deadKeysOpt.default(SortedSet[]);
+    refCountsOpt: ?KeyMap<Int>,
+    newKeysOpt: ?KeySet,
+    deadKeysOpt: ?KeySet,
+  ): (?KeyMap<Int>, ?this) {
+    refCounts = this.castOrEmptyKeyMap(refCountsOpt);
+    newKeys = this.castOrEmptyKeySet(newKeysOpt);
+    deadKeys = this.castOrEmptyKeySet(deadKeysOpt);
 
     toRemove = mutable Vector[];
     !refCounts = refCounts.mergeWith2(
-      newKeys.inner,
-      deadKeys.inner,
+      newKeys.getInner(),
+      deadKeys.getInner(),
       (key, refCountOpt, newOpt, deadOpt) -> {
         (refCountOpt, newOpt, deadOpt) match {
         | (_, None(), None())
@@ -108,13 +108,13 @@ class TLazyDir{
 
   fun update(
     context: mutable Context,
-    dirtyReadersOpt: ?SortedMap<DirName, SortedSet<Key>>,
+    dirtyReadersOpt: ?SortedMap<DirName, KeySet>,
   ): void {
     dirtyReadersOpt match {
     | None() -> void
     | Some(dirtyReaders) ->
       for (dirtyReaderDirName => dirtyReaderKeys in dirtyReaders) {
-        for (key in dirtyReaderKeys) {
+        for (key in this.castKeySet(dirtyReaderKeys)) {
           oldValue = this.data.maybeGet(key);
           !this.data = this.data.remove(key);
           context.updateDirtyReaders(Path::create(dirtyReaderDirName, key));

skiplang/prelude/src/skstore/Path.sk

@@ -247,6 +247,59 @@ class SID(value: String) extends Key {
   }
 }
 
+base class KeySet uses Unsafe.Downcastable
+
+class TKeySet<+K: Orderable> private (inner: SortedSet<K>) extends KeySet {
+  static fun empty(): this {
+    static(SortedSet[])
+  }
+
+  fun set<K2: Orderable>[K: K2](key: K2): TKeySet<K2> {
+    TKeySet(this.inner.set(key))
+  }
+
+  fun values(): mutable Iterator<K> {
+    this.inner.values()
+  }
+
+  fun getInner(): TKeyMap<K, void> {
+    TKeyMap<K, void>::ofKeySet(this)
+  }
+}
+
+base class KeyMap<+V> uses Unsafe.Downcastable
+
+class TKeyMap<+K: Orderable, +V: frozen> private (
+  inner: SortedMap<K, V>,
+) extends KeyMap<V> {
+  static fun empty(): this {
+    static(SortedMap[])
+  }
+
+  static fun ofKeySet[V: void](set: TKeySet<K>): TKeyMap<K, void> {
+    TKeyMap<K, void>(set.inner.inner)
+  }
+
+  fun isEmpty(): Bool {
+    this.inner.isEmpty()
+  }
+
+  fun set<K2: Orderable, V2: frozen>[K: K2, V: V2](
+    key: K2,
+    value: V2,
+  ): TKeyMap<K2, V2> {
+    TKeyMap(this.inner.set(key, value))
+  }
+
+  fun mergeWith2<K2: Orderable, U: frozen, W: frozen, R: frozen>[K: K2](
+    other1: TKeyMap<K2, U>,
+    other2: TKeyMap<K2, W>,
+    f: (K2, ?V, ?U, ?W) -> ?R,
+  ): TKeyMap<K2, R> {
+    TKeyMap(this.inner.mergeWith2(other1.inner, other2.inner, f))
+  }
+}
+
 base class Exception extends .Exception uses Show
 
 fun error<T>(msg: String): T {