DMap.sk

/******************************************************************************
 * Efficient Map for "diffing". The server will need to be able to quickly
 * respond when a client reconnects. Typically, we will need to be able to
 * quickly send the files that have changed in a directory since the last
 * connection. We achieve this by introducing an AVL tree that also
 * maintains the youngest element present in the subtree.
 * This way, we can walk the tree and prune all the parts that are "too old".
 ******************************************************************************/

module SKFS;

value class TimeTag(max: Int, current: Int) {
  static fun create(time: Int): this {
    TimeTag(time, time)
  }
}

mutable class IterHolder<Key, Value>(iter: mutable Iterator<(Key, Value)>) {
  mutable fun items(): mutable Iterator<(Key, Value)> {
    this.iter
  }
}

base class Boundary<T> {
  children =
  | Inclusive(T)
  | Exclusive(T)
}

base class DMap<Key: FileName, Value> {
  children =
  | Nil()
  | Node{
    time: TimeTag,
    height: Int,
    key: Key,
    value: Value,
    left: DMap<Key, Value>,
    right: DMap<Key, Value>,
  }

  fun isEmpty(): Bool {
    this is Nil _
  }

  fun items(): mutable Iterator<(Key, Value)> {
    this match {
    | Nil() -> void
    | Node{key, value, left, right} ->
      for (x in left.items()) {
        yield x;
      };
      yield (key, value);
      for (x in right.items()) {
        yield x;
      }
    }
  }

  private fun itemsAfterKeyHelper(
    boundary: Boundary<Key>,
  ): mutable Iterator<(Key, Value)> {
    this match {
    | Nil() -> void
    | Node{key, value, left, right} ->
      if (
        boundary match {
        | Inclusive(lastSkipped) -> key >= lastSkipped
        | Exclusive(lastSkipped) -> key > lastSkipped
        }
      ) {
        for (x in left.itemsAfterKeyHelper(boundary)) {
          yield x;
        };
        yield (key, value);
      };
      for (x in right.itemsAfterKeyHelper(boundary)) {
        yield x;
      }
    }
  }

  fun itemsAfterKey(
    lastSkippedOpt: ?Boundary<Key>,
  ): mutable IterHolder<Key, Value> {
    mutable IterHolder(
      lastSkippedOpt match {
      | None() -> this.items()
      | Some(lastSkipped) -> this.itemsAfterKeyHelper(lastSkipped)
      },
    )
  }

  fun minKey(cond: ((Key, Value)) -> Bool): ?Key {
    this.items().filter(cond).next().map(x -> x.i0)
  }

  fun itemsWithTime(): mutable Iterator<(Key, Value, TimeTag)> {
    this match {
    | Nil() -> void
    | Node{time, key, value, left, right} ->
      for (x in left.itemsWithTime()) {
        yield x;
      };
      yield (key, value, time);
      for (x in right.itemsWithTime()) {
        yield x;
      }
    }
  }

  fun eachWithTime(f: (Key, Int) -> void): void
  | Nil() -> void
  | Node{key, time, left, right} ->
    left.eachWithTime(f);
    f(key, time.current);
    right.eachWithTime(f)

  fun getMaxTime(): Int
  | Nil() -> 0
  | Node{time} -> time.max

  fun getHeight(): Int
  | Nil() -> 0
  | Node{height} -> height

  static fun empty<Key: FileName, Value>(): DMap<Key, Value> {
    Nil()
  }

  fun getChangesAcc(after: Int, acc: mutable Set<Key>): void
  | Nil() -> void
  | Node{time, key, left, right} ->
    if (time.max < after) return void;
    if (time.current >= after) {
      acc.insert(key);
    };
    left.getChangesAcc(after, acc);
    right.getChangesAcc(after, acc)

  fun getChangesAfter(time: Int): Set<Key> {
    acc = mutable Set[];
    this.getChangesAcc(time, acc);
    unsafe_chill_trust_me(acc)
  }

  fun getChangesRangeAcc(
    timeStart: Int,
    timeEnd: Int,
    acc: mutable Vector<(Key, Value)>,
  ): void
  | Nil() -> void
  | Node{time, key, value, left, right} ->
    if (time.max < timeStart) return void;
    if (time.current >= timeStart && time.current < timeEnd) {
      acc.push((key, value));
    };
    left.getChangesRangeAcc(timeStart, timeEnd, acc);
    right.getChangesRangeAcc(timeStart, timeEnd, acc)

  fun getChangesRange(timeStart: Int, timeEnd: Int): Vector<(Key, Value)> {
    acc = mutable Vector[];
    this.getChangesRangeAcc(timeStart, timeEnd, acc);
    unsafe_chill_trust_me(acc)
  }

  static fun node(
    time: TimeTag,
    key: Key,
    value: Value,
    left: DMap<Key, Value>,
    right: DMap<Key, Value>,
  ): DMap<Key, Value> {
    !time.max = max(time.max, max(left.getMaxTime(), right.getMaxTime()));
    height = 1 + max(left.getHeight(), right.getHeight());
    Node{time, height, key, value, left, right}
  }

  fun set(time: TimeTag, key: Key, value: Value): DMap<Key, Value>
  | Nil() -> static::node(time, key, value, Nil(), Nil())
  | Node{time => t, left => l, right => r, key => k, value => v} ->
    !t.max = max(time.max, t.max);
    !time.max = t.max;
    compare(key, this.key) match {
    | LT() -> static::balance(t, k, v, l.set(time, key, value), r)
    | EQ() -> static::node(time, k, value, l, r)
    | GT() -> static::balance(t, k, v, l, r.set(time, key, value))
    }

  fun containsKey(k: Key): Bool
  | Nil() -> false
  | Node{key, left, right} ->
    k.compare(key) match {
    | LT() -> left.containsKey(k)
    | EQ() -> true
    | GT() -> right.containsKey(k)
    }

  fun maybeGet(k: Key): ?Value
  | Nil() -> None()
  | Node{key, value, left, right} ->
    k.compare(key) match {
    | LT() -> left.maybeGet(k)
    | EQ() -> Some(value)
    | GT() -> right.maybeGet(k)
    }

  fun remove(k: Key): DMap<Key, Value>
  | Nil() -> this
  | Node{time, key, value, left, right} ->
    k.compare(key) match {
    | LT() -> static::node(time, key, value, left.remove(k), right)
    | EQ() -> DMap::merge(left, right)
    | GT() -> static::node(time, key, value, left, right.remove(k))
    }

  fun removeMinBinding(): ((TimeTag, Key, Value), DMap<Key, Value>) {
    this match {
    | Nil() -> invariant_violation("Remove min binding on empty DMap")
    | Node{time, left => Nil(), key, value, right} ->
      ((time, key, value), right)
    | Node{time, left, key, value, right} ->
      (binding, !left) = left.removeMinBinding();
      (binding, static::balance(time, key, value, left, right))
    }
  }

  static fun merge(
    t1: DMap<Key, Value>,
    t2: DMap<Key, Value>,
  ): DMap<Key, Value> {
    (t1, t2) match {
    | (Nil(), t) -> t
    | (t, Nil()) -> t
    | (_, _) ->
      ((t, k, v), !t2) = t2.removeMinBinding();
      static::balance(t, k, v, t1, t2)
    }
  }

  protected static fun balance(
    t: TimeTag,
    k: Key,
    v: Value,
    l: DMap<Key, Value>,
    r: DMap<Key, Value>,
  ): DMap<Key, Value> {
    hl = l.getHeight();
    hr = r.getHeight();
    if (hl > hr + 2) {
      l match {
      | Nil() -> invariant_violation("DMap: empty left tree")
      | Node{time => lt, left => ll, right => lr, key => lk, value => lv} ->
        if (ll.getHeight() >= lr.getHeight()) {
          static::node(lt, lk, lv, ll, static::node(t, k, v, lr, r))
        } else {
          lr match {
          | Nil() -> invariant_violation("DMap: empty l.right tree")
          | Node{
            time => lrt,
            left => lrl,
            right => lrr,
            key => lrk,
            value => lrv,
          } ->
            static::node(
              lrt,
              lrk,
              lrv,
              static::node(lt, lk, lv, ll, lrl),
              static::node(t, k, v, lrr, r),
            )
          }
        }
      }
    } else if (hr > hl + 2) {
      r match {
      | Nil() -> invariant_violation("DMap empty right tree")
      | Node{time => rt, left => rl, right => rr, key => rk, value => rv} ->
        if (rr.getHeight() >= rl.getHeight()) {
          static::node(rt, rk, rv, static::node(t, k, v, l, rl), rr)
        } else {
          rl match {
          | Nil() -> invariant_violation("DMap empty r.left tree")
          | Node{
            time => rlt,
            left => rll,
            right => rlr,
            key => rlk,
            value => rlv,
          } ->
            static::node(
              rlt,
              rlk,
              rlv,
              static::node(t, k, v, l, rll),
              static::node(rt, rk, rv, rlr, rr),
            )
          }
        }
      }
    } else {
      static::node(t, k, v, l, r)
    }
  }
}

module end;

module SKFSTest;

fun testDMap(): void {
  x = SKFS.DMap::empty();
  !x = x.set(SKFS.TimeTag::create(1), SKFS.IID(1), 1);
  !x = x.set(SKFS.TimeTag::create(1), SKFS.IID(2), 2);
  !x = x.set(SKFS.TimeTag::create(1), SKFS.IID(3), 3);
  !x = x.set(SKFS.TimeTag::create(2), SKFS.IID(1), 23);
  expectEq("Changes after", Array[1, 2, 3], () ->
    x.getChangesAfter(1).toArray().map(x -> x.hashTag)
  );
  void
}

module end;