Adding in std/data/sorted-dict.kk

lib/std/data/red-black-tree.kk

@@ -91,9 +91,9 @@ public fun to-list(tree: tree<a>): list<a> {
 
 
 // TODO should this be called unsafe ?
-public fun from-list(list: list<a>, compare: (a, a) -> e order): e tree<a> {
+public fun from-list(list: list<a>, map: (a) -> e b, compare: (b, a) -> e order): e tree<a> {
   list.foldl(empty) fun(l, r) {
-    unsafe-insert(l, r, compare, True)
+    unsafe-insert(l, map(r), r, compare, True)
   }
 }
 
@@ -202,13 +202,13 @@ fun redden(tree: tree<a>): tree<a> {
 }
 
 
-fun unsafe-insert1(tree: tree<a>, value: a, compare: (a, a) -> e order, replace: bool): e tree<a> {
+fun unsafe-insert1(tree: tree<a>, key: b, value: a, compare: (b, a) -> e order, replace: bool): e tree<a> {
   match (tree) {
     Node(left, middle, right, color) ->
-      match (compare(value, middle)) {
-        Lt -> balance(Node(unsafe-insert1(left, value, compare, replace), middle, right, color))
+      match (compare(key, middle)) {
+        Lt -> balance(Node(unsafe-insert1(left, key, value, compare, replace), middle, right, color))
 
-        Gt -> balance(Node(left, middle, unsafe-insert1(right, value, compare, replace), color))
+        Gt -> balance(Node(left, middle, unsafe-insert1(right, key, value, compare, replace), color))
 
         Eq -> if (replace) {
           // TODO in this situation we should avoid balancing, because the structure hasn't changed
@@ -224,25 +224,25 @@ fun unsafe-insert1(tree: tree<a>, value: a, compare: (a, a) -> e order, replace:
 }
 
 // This is unsafe because it's possible to use different compare functions on the same tree
-public fun unsafe-insert(tree: tree<a>, value: a, compare: (a, a) -> e order, replace: bool): e tree<a> {
-  blacken(unsafe-insert1(tree, value, compare, replace))
+public fun unsafe-insert(tree: tree<a>, key: b, value: a, compare: (b, a) -> e order, replace: bool): e tree<a> {
+  blacken(unsafe-insert1(tree, key, value, compare, replace))
 }
 
 
-fun unsafe-remove1(tree: tree<a>, value: a, compare: (a, a) -> e order): e tree<a> {
+fun unsafe-remove1(tree: tree<a>, key: b, compare: (b, a) -> e order): e tree<a> {
   match (tree) {
     Node(left, middle, right, _) ->
-      match (compare(value, middle)) {
+      match (compare(key, middle)) {
         Lt -> match (left) {
-          Node(_, _, _, Black) -> balance-left(unsafe-remove1(left, value, compare), middle, right)
-          _ -> Node(unsafe-remove1(left, value, compare), middle, right, Red)
+          Node(_, _, _, Black) -> balance-left(unsafe-remove1(left, key, compare), middle, right)
+          _ -> Node(unsafe-remove1(left, key, compare), middle, right, Red)
         }
 
         Eq -> unsafe-merge(left, right)
 
         Gt -> match (right) {
-          Node(_, _, _, Black) -> balance-right(left, middle, unsafe-remove1(right, value, compare))
-          _ -> Node(left, middle, unsafe-remove1(right, value, compare), Red)
+          Node(_, _, _, Black) -> balance-right(left, middle, unsafe-remove1(right, key, compare))
+          _ -> Node(left, middle, unsafe-remove1(right, key, compare), Red)
         }
       }
 
@@ -251,6 +251,6 @@ fun unsafe-remove1(tree: tree<a>, value: a, compare: (a, a) -> e order): e tree<
 }
 
 // This is unsafe because it's possible to use different compare functions on the same tree
-public fun unsafe-remove(tree: tree<a>, value: a, compare: (a, a) -> e order): e tree<a> {
-  blacken(unsafe-remove1(tree, value, compare))
+public fun unsafe-remove(tree: tree<a>, key: b, compare: (b, a) -> e order): e tree<a> {
+  blacken(unsafe-remove1(tree, key, compare))
 }

lib/std/data/sorted-dict.kk

@@ -0,0 +1,94 @@
+/*---------------------------------------------------------------------------
+  Copyright 2017 Microsoft Corporation.
+
+  This is free software; you can redistribute it and/or modify it under the
+  terms of the Apache License, Version 2.0. A copy of the License can be
+  found in the file "license.txt" at the root of this distribution.
+---------------------------------------------------------------------------*/
+
+// TODO unit tests
+// TODO maybe allow for compare functions with effects
+// TODO maybe add in the other functions from sorted-set
+module std/data/sorted-dict
+
+import red-black-tree
+
+
+abstract struct sorted-dict<k, a>(
+  compare: (k, (k, a)) -> order,
+  tree: tree<(k, a)>
+)
+
+fun from-compare(compare: (k, k) -> order): ((k, (k, a)) -> order) {
+  (fun(k1, v2) {
+    match (v2) {
+      (k2, _) -> compare(k1, k2)
+    }
+  })
+}
+
+public fun from-list(list: list<(k, a)>, compare: (k, k) -> order): sorted-dict<k, a> {
+  val compare2 = from-compare(compare)
+  Sorted-dict(compare2, from-list(list, fst, compare2))
+}
+
+public fun to-list(dict: sorted-dict<k, a>): list<(k, a)> {
+  to-list(dict.tree)
+}
+
+public fun empty(compare: (k, k) -> order): sorted-dict<k, a> {
+  Sorted-dict(from-compare(compare), empty)
+}
+
+public fun get(dict: sorted-dict<k, a>, key: k): maybe<a> {
+  unsafe-lookup(dict.tree, key, dict.compare).map(snd)
+}
+
+public fun has?(dict: sorted-dict<k, a>, key: k): bool {
+  // TODO this can be implemented faster
+  bool(get(dict, key))
+}
+
+public fun set(dict: sorted-dict<k, a>, key: k, value: a): sorted-dict<k, a> {
+  dict(tree = unsafe-insert(dict.tree, key, (key, value), dict.compare, True))
+}
+
+public fun remove(dict: sorted-dict<k, a>, key: k): sorted-dict<k, a> {
+  dict(tree = unsafe-remove(dict.tree, key, dict.compare))
+}
+
+public fun merge-left(left: sorted-dict<k, a>, right: sorted-dict<k, a>): sorted-dict<k, a> {
+  val compare = left.compare
+  left(
+    tree = right.tree.foldl(left.tree) fun(old, a) {
+      match (a) {
+        // If a value exists in both `left` and `right` it will prefer the value from `left`
+        (key, _) -> unsafe-insert(old, key, a, compare, False)
+      }
+    }
+  )
+}
+
+public fun merge-right(left: sorted-dict<k, a>, right: sorted-dict<k, a>): sorted-dict<k, a> {
+  val compare = right.compare
+  right(
+    tree = left.tree.foldl(right.tree) fun(old, a) {
+      match (a) {
+        // If a value exists in both `left` and `right` it will prefer the value from `right`
+        (key, _) -> unsafe-insert(old, key, a, compare, False)
+      }
+    }
+  )
+}
+
+public fun subset-of?(smaller: sorted-dict<k, a>, bigger: sorted-dict<k, a>): bool {
+  smaller.tree.find-first fun(a) {
+    match (a) {
+      (key, _) -> if (bigger.has?(key)) {
+        Nothing
+      } else {
+        Just(False)
+      }
+    }
+  }.default(True)
+}

lib/std/data/sorted-set.kk

@@ -22,7 +22,7 @@ abstract struct sorted-set<a>(
 )
 
 public fun from-list(list: list<a>, compare: (a, a) -> order): sorted-set<a> {
-  Sorted-set(compare, from-list(list, compare))
+  Sorted-set(compare, from-list(list, id, compare))
 }
 
 public fun to-list(set: sorted-set<a>): list<a> {
@@ -39,7 +39,7 @@ public fun has?(set: sorted-set<a>, value: a): bool {
 }
 
 public fun insert(set: sorted-set<a>, value: a): sorted-set<a> {
-  set(tree = unsafe-insert(set.tree, value, set.compare, True))
+  set(tree = unsafe-insert(set.tree, value, value, set.compare, True))
 }
 
 public fun remove(set: sorted-set<a>, value: a): sorted-set<a> {
@@ -60,7 +60,7 @@ public fun intersection-right(left: sorted-set<a>, right: sorted-set<a>): sorted
   right(
     tree = right.tree.foldl(empty) fun(old, value) {
       if (left.has?(value)) {
-        unsafe-insert(old, value, compare, True)
+        unsafe-insert(old, value, value, compare, True)
       } else {
         old
       }
@@ -87,7 +87,7 @@ public fun union-left(left: sorted-set<a>, right: sorted-set<a>): sorted-set<a>
   left(
     tree = right.tree.foldl(left.tree) fun(old, value) {
       // If a value exists in both `left` and `right` it will prefer the value from `left`
-      unsafe-insert(old, value, compare, False)
+      unsafe-insert(old, value, value, compare, False)
     }
   )
 }
@@ -97,7 +97,7 @@ public fun union-right(left: sorted-set<a>, right: sorted-set<a>): sorted-set<a>
   right(
     tree = left.tree.foldl(right.tree) fun(old, value) {
       // If a value exists in both `left` and `right` it will prefer the value from `right`
-      unsafe-insert(old, value, compare, False)
+      unsafe-insert(old, value, value, compare, False)
     }
   )
 }