Basic dictionaries

examples/stdlib/acme.effekt

@@ -9,6 +9,7 @@ import buffer
 import bytearray
 import char
 import dequeue
+import dictionary
 import effekt
 import exception
 import heap

examples/stdlib/dictionary/each_tagged.check

@@ -0,0 +1,6 @@
+a
+apple
+ab
+amazing
+b
+banana

examples/stdlib/dictionary/each_tagged.effekt

@@ -0,0 +1,16 @@
+import dictionary
+import stream
+import bytearray
+
+def main() = {
+  var trie = Empty()
+
+  trie = trie.insert("a", "apple")
+  trie = trie.insert("ab", "amazing")
+  trie = trie.insert("b", "banana")
+
+  for[Tagged[String]] { eachTagged(trie) } { case Tagged(key, value) =>
+    println(bytearray::toString(key))
+    println(value)
+  }
+}

examples/stdlib/dictionary/lookup_insert.check

@@ -0,0 +1,4 @@
+Some(1)
+Some(2)
+Some(3)
+None()

examples/stdlib/dictionary/lookup_insert.effekt

@@ -0,0 +1,14 @@
+import dictionary
+
+def main() = {
+  var trie = Empty()
+
+  trie = trie.insert("hello", 1)
+  trie = trie.insert("world", 2)
+  trie = trie.insert("help", 3)
+
+  println(trie.lookup("hello"))
+  println(trie.lookup("world"))
+  println(trie.lookup("help"))
+  println(trie.lookup("missing"))
+}

libraries/common/dictionary.effekt

@@ -0,0 +1,133 @@
+module dictionary
+
+import array
+import bytearray
+import stream
+
+
+/// A trie-based dictionary for efficient key-value storage with bytearray keys.
+/// Supports String and ByteArray keys with fast lookup and insertion.
+type Dictionary[A] {
+  Empty()
+  Leaf(value: A)
+  Branch(children: Array[Dictionary[A]])
+  Node(value: A, children: Array[Dictionary[A]])
+}
+
+/// Insert a value with a String key.
+def insert[A](trie: Dictionary[A], key: String, value: A): Dictionary[A] =
+  insert(trie, key.fromString, value)
+
+/// Insert a key-value pair into the dictionary with a ByteArray key.
+def insert[A](trie: Dictionary[A], key: ByteArray, value: A): Dictionary[A] = {
+  def go(trie: Dictionary[A], i: Int): Dictionary[A] =
+    if (i < key.size) {
+      trie match {
+        case Empty() =>
+          val child = go(Empty(), i + 1)
+          val children = array(256, Empty())
+          children.unsafeSet(key.unsafeGet(i).toInt, child)
+          Branch(children)
+        case Leaf(existing) =>
+          val child = go(Empty(), i + 1)
+          val children = array(256, Empty())
+          children.unsafeSet(key.unsafeGet(i).toInt, child)
+          Node(existing, children)
+        case Branch(children) =>
+          val child = go(children.unsafeGet(key.unsafeGet(i).toInt), i + 1)
+          children.unsafeSet(key.unsafeGet(i).toInt, child)
+          Branch(children)
+        case Node(existing, children) =>
+          val child = go(children.unsafeGet(key.unsafeGet(i).toInt), i + 1)
+          children.unsafeSet(key.unsafeGet(i).toInt, child)
+          Node(existing, children)
+      }
+    } else {
+      trie match {
+        case Empty() => Leaf(value)
+        case Leaf(_) => Leaf(value)
+        case Node(_, children) => Node(value, children)
+        case Branch(children) => Node(value, children)
+      }
+    }
+  go(trie, 0)
+}
+
+/// Look up a value in the dictionary by string key.
+def lookup[A](trie: Dictionary[A], key: String): Option[A] =
+  lookup(trie, key.fromString)
+
+/// Look up a value in the dictionary by bytearray key.
+def lookup[A](trie: Dictionary[A], key: ByteArray): Option[A] = {
+  def go(trie: Dictionary[A], i: Int): Option[A] =
+    if (i < key.size) {
+      trie match {
+        case Empty() => None()
+        case Leaf(_) => None()
+        case Branch(children) => go(children.unsafeGet(key.unsafeGet(i).toInt), i + 1)
+        case Node(_, children) => go(children.unsafeGet(key.unsafeGet(i).toInt), i + 1)
+      }
+    } else {
+      trie match {
+        case Empty() => None()
+        case Leaf(value) => Some(value)
+        case Branch(_) => None()
+        case Node(value, _) => Some(value)
+      }
+    }
+  go(trie, 0)
+}
+
+/// A record pairing a key with a value for use in streams.
+record Tagged[A](key: ByteArray, value: A)
+
+/// Collect key-value pairs from a stream of Tagged elements into a Dictionary.
+def collectTrie[A] { stream: () => Unit / emit[Tagged[A]] }: Dictionary[A] = {
+  var trie = Empty()
+  for[Tagged[A]] { stream() } { case Tagged(bytes, value) =>
+    trie = trie.insert(bytes, value)
+  }
+  return trie
+}
+
+/// Emit all values stored in the dictionary as a push stream.
+def each[A](trie: Dictionary[A]): Unit / emit[A] =
+  trie match {
+    case Empty() => ()
+    case Leaf(value) => do emit(value)
+    case Branch(children) =>
+      each(0, children.size) { i =>
+        each(children.unsafeGet(i))
+      }
+    case Node(value, children) =>
+      do emit(value)
+      each(0, children.size) { i =>
+        each(children.unsafeGet(i))
+      }
+  }
+
+/// Emit all key-value pairs stored in the dictionary as a push stream of Tagged elements.
+def eachTagged[A](trie: Dictionary[A]): Unit / emit[Tagged[A]] = {
+  def fromBytesReverse(bytes: List[Byte], n: Int): ByteArray = {
+    val key = bytearray::allocate(n)
+    foreachIndex(bytes) { (i, b) => key.unsafeSet(n - 1 - i, b) }
+    key
+  }
+  def go(trie: Dictionary[A], key: List[Byte], n: Int): Unit =
+    trie match {
+      case Empty() => ()
+      case Leaf(value) => do emit(Tagged(fromBytesReverse(key, n), value))
+      case Branch(children) =>
+        each(0, children.size) { i =>
+          go(children.unsafeGet(i), Cons(i.toByte, key), n + 1)
+        }
+      case Node(value, children) =>
+        do emit(Tagged(fromBytesReverse(key, n), value))
+        each(0, children.size) { i =>
+          go(children.unsafeGet(i), Cons(i.toByte, key), n + 1)
+        }
+    }
+  go(trie, Nil(), 0)
+}
+
+