use projection to get Set and Dict

example/dictionary.sw

@@ -3,26 +3,29 @@ def black = false end
 def red = true end
 // TODO: s/\((.*\.c) == red\)/$1.red/g
 
-tydef RBTree k v = rec b. Unit + [c: Color, k: k, v: v, l: b, r: b] end
-tydef RBNode k v = rec n. [c: Color, k: k, v: v, l: Unit + n, r: Unit + n] end
+tydef Maybe a = Unit + a end
+def mmap : (a -> b) -> Maybe a -> Maybe b = \f.\m. case m (\_. inl ()) (\a. inr (f a)) end
 
-def emptyT : RBTree k v = inl () end
+tydef RBTree k = rec b. Unit + [c: Color, k: k, l: b, r: b] end
+tydef RBNode k = rec n. [c: Color, k: k, l: Unit + n, r: Unit + n] end
 
-def isEmptyT : RBTree k v -> Bool = \d. d == emptyT end
+def emptyT : RBTree k = inl () end
 
-def getT : k -> RBTree k v -> Unit + v = \x.\t.
+def isEmptyT : RBTree k -> Bool = \d. d == emptyT end
+
+def getT : (k -> a) -> a -> RBTree k -> Unit + k = \p.\x.\t.
   case t (\_. inl ()) (\n.
-    if (x == n.k) { inr n.v } {
-        if (x < n.k) {
-            getT x n.l
+    if (x == p n.k) { inr n.k } {
+        if (x < p n.k) {
+            getT p x n.l
         } {
-            getT x n.r
+            getT p x n.r
         }
     }
   )
 end
 
-def containsT : k -> RBTree k v -> Bool = \x.\t. getT x t != inl () end
+def containsT : (k -> a) -> a -> RBTree k -> Bool = \p.\x.\t. getT p x t != inl () end
 
 tydef List a = rec l. Unit + (a * l) end
 
@@ -39,9 +42,9 @@ def formatL : List a -> Text = \l.
  in case l (\_. "[]") (\n. "[" ++ format (fst n) ++ f (snd n))
 end
 
-def inorder : RBTree k v -> List (k * v) = \t.
+def inorder : RBTree k -> List k = \t.
   case t (\_. inl ()) (\n.
-    appendL (inorder n.l) (appendL (pureL (n.k, n.v)) $ inorder n.r)
+    appendL (inorder n.l) (appendL (pureL n.k) $ inorder n.r)
   )
 end
 
@@ -52,32 +55,29 @@ balance B a x0 (T R (T R b x1 c) x2 d) = T R (T B a x0 b) x1 (T B c x2 d) -- cas
 balance B a x0 (T R b x1 (T R c x2 d)) = T R (T B a x0 b) x1 (T B c x2 d) -- case 4: RR red
 balance color a x b = T color a x b -- case 5
 */
-def balanceT : RBNode k v -> RBTree k v = \t.
-  let balanced : k -> v -> RBTree k v -> RBTree k v
-              -> k -> v
-              -> k -> v -> RBTree k v -> RBTree k v
-              -> RBTree k v
-    = \x0.\y0.\a.\b. \x1.\y1. \x2.\y2.\c.\d.
-    inr [c=red, k=x1, v=y1, l=inr [c=black, k=x0, v=y0, l=a, r=b], r=inr [c=black, k=x2, v=y2, l=c, r=d]]
-  in let case5 : RBNode k v -> RBTree k v =
+def balanceT : RBNode k -> RBTree k = \t.
+  let balanced : [ a: RBTree k, x0: k, b: RBTree k,   x1: k,   c: RBTree k, x2: k, d: RBTree k] -> RBTree k
+    = \v.
+    inr [c=red, l=inr [c=black, l=v.a, k=v.x0, r=v.b], k=v.x1, r=inr [c=black, l=v.c, k=v.x2, r=v.d]]
+  in let case5 : RBNode k -> RBTree k =
     inr  // just id - TODO: inline?
-  in let balanceRR : RBNode k v -> RBNode k v -> RBTree k v = \t.\rn.
+  in let balanceRR : RBNode k -> RBNode k -> RBTree k = \t.\rn.
     case rn.r (\_. case5 t) (\rrn.
       if (rrn.c == red) {
-        balanced t.k t.v t.l rn.l  rn.k rn.v  rrn.k rrn.v rrn.l rrn.r
+        balanced [a=t.l, x0=t.k, b=rn.l, x1=rn.k, c=rrn.l, x2=rrn.k, d=rrn.r]
       } {
         case5 t
       }
     )
-  in let balanceRL : RBNode k v -> RBNode k v -> RBTree k v = \t.\rn.
+  in let balanceRL : RBNode k -> RBNode k -> RBTree k = \t.\rn.
     case rn.l (\_. balanceRR t rn) (\rln.
       if (rln.c == red) {
-        balanced t.k t.v t.l rln.r  rln.k rln.v  rn.k rn.v rln.r rn.r
+        balanced [a=t.l, x0=t.k, b=rln.r, x1=rln.k, c=rln.r, x2=rn.k, d=rn.r]
       } {
         balanceRR t rn
       }
     )
-  in let balanceR : RBNode k v -> RBTree k v = \t.
+  in let balanceR : RBNode k -> RBTree k = \t.
     case t.r (\_.
       case5 t
     ) (\lr.
@@ -87,23 +87,23 @@ def balanceT : RBNode k v -> RBTree k v = \t.
         case5 t
       }
     )
-  in let balanceLR : RBNode k v -> RBNode k v -> RBTree k v = \t.\ln.
+  in let balanceLR : RBNode k -> RBNode k -> RBTree k = \t.\ln.
     case ln.r (\_. balanceR t) (\lrn.
       if (lrn.c == red) {
-        balanced ln.k ln.v ln.l lrn.l  lrn.k lrn.v  t.k t.v lrn.l t.r
+        balanced [a=ln.l, x0=ln.k, b=lrn.l, x1=lrn.k, c=lrn.l, x2=t.k, d=t.r]
       } {
         balanceR t
       }
     )
-  in let balanceLL : RBNode k v -> RBNode k v -> RBTree k v = \t.\ln.
+  in let balanceLL : RBNode k -> RBNode k -> RBTree k = \t.\ln.
     case ln.l (\_. balanceLR t ln) (\lln.
       if (lln.c == red) {
-        balanced lln.k lln.v lln.l lln.r  ln.k ln.v  t.k t.v ln.r t.r
+        balanced [a=lln.l, x0=lln.k, b=lln.r, x1=ln.k, c=ln.r, x2=t.k, d=t.r]
       } {
         balanceLR t ln
       }
     )
-  in let balanceL : RBNode k v -> RBTree k v = \t.
+  in let balanceL : RBNode k -> RBTree k = \t.
     case t.l (\_.
       balanceR t
     ) (\ln.
@@ -120,25 +120,25 @@ def balanceT : RBNode k v -> RBTree k v = \t.
   }
 end
 
-def insertT : k -> v -> RBTree k v -> RBTree k v = \x.\y.\t.
-  let ins : k -> v -> RBTree k v -> RBTree k v  = \x.\y.\t. case t
+def insertT : (k -> a) -> k -> RBTree k -> RBTree k = \p.\x.\t.
+  let ins : (k -> a) -> k -> RBTree k -> RBTree k =\p.\x.\t. case t
     (\_.
-      inr [c=red, k=x, v=y, l=inl (), r=inl ()]
+      inr [c=red, k=x, l=inl (), r=inl ()]
     ) (\n.
-      if (x == n.k) { inr [c=n.c, k=x, v=y, l=n.l, r=n.r] } {
+      if (p x == p n.k) { inr [c=n.c, k=x, l=n.l, r=n.r] } {
         if (x < n.k) {
-            balanceT [c=n.c, k=n.k, v=n.v, l=ins x y n.l, r=n.r]
+            balanceT [c=n.c, k=n.k, l=ins p x n.l, r=n.r]
         } {
-            balanceT [c=n.c,k=n.k, v=n.v, l=n.l, r=ins x y n.r]
+            balanceT [c=n.c,k=n.k, l=n.l, r=ins p x n.r]
         }
       }
     )
-  in let makeBlack : RBTree k v -> RBTree k v = \t. case t (\_.
+  in let makeBlack : RBTree k -> RBTree k = \t. case t (\_.
         fail "makeBlack will always be called on nonempty"
       ) (\n.
-        inr [c=black, k=n.k, v=n.v, l=n.l, r=n.r]
+        inr [c=black, k=n.k, l=n.l, r=n.r]
       )
-  in makeBlack $ ins x y t
+  in makeBlack $ ins p x t
 end
 
 /*
@@ -171,9 +171,33 @@ def deleteD : k -> RBTree k v -> RBTree k v = \x.\d.
 end
 */
 
+/*******************************************************************/
+/*                           DICTIONARY                            */
+/*******************************************************************/
+
+tydef Dict k v = RBTree (k * v) end
+
+def emptyD : Dict k v = emptyT end
+def isEmptyD : Dict k v -> Bool = isEmptyT end
+def getD : k -> Dict k v -> Unit + v = \k.\d. mmap snd (getT fst k d) end 
+def containsD : k -> Dict k v -> Bool = containsT fst end
+def insertD : k -> v -> Dict k v -> Dict k v = \k.\v. insertT fst (k, v) end
+
+/*******************************************************************/
+/*                              SET                                */
+/*******************************************************************/
+
+tydef Set k = RBTree k end
+
+def emptyS : Set k = emptyT end
+def isEmptyS : Set k -> Bool = isEmptyT end
+def containsS : k -> Set k -> Bool = containsT (\x.x) end
+def insertS : k -> Set k -> Set k = insertT (\x.x) end
+
 /*******************************************************************/
 /*                           UNIT TESTS                            */
 /*******************************************************************/
+
 def indent = \i. if (i <= 0) {""} {"  " ++ indent (i - 1)} end
 
 def assert = \i. \b. \m. if b {log $ indent i ++ "OK: " ++ m} {log "FAIL:"; fail m} end
@@ -194,25 +218,25 @@ end
 
 def test_empty: Int -> Cmd Unit = \i.
   group i "EMPTY TESTS" (\i.
-    assert i (isEmptyT emptyT) "empty tree should be empty";
-    assert_eq i (inl ()) (getT 0 $ emptyT) "no element should be present in an empty tree";
-    assert i (not $ containsT 0 $ emptyT) "empty tree should not contain any elements";
+    assert i (isEmptyD emptyD) "empty tree should be empty";
+    assert_eq i (inl ()) (getD 0 $ emptyD) "no element should be present in an empty tree";
+    assert i (not $ containsD 0 $ emptyD) "empty tree should not contain any elements";
   )
 end
 
 def test_insert: Int -> Cmd Unit = \i.
   group i "INSERT TESTS" (\i.
     group i "test {0: 1}" (\i.
-      let tree0 = insertT 0 1 emptyT in
-      assert i (not $ isEmptyT tree0) "nonempty tree should not be empty";
-      assert_eq i (inr 1) (getT 0 $ tree0) "one element should be present in a one element tree";
-      assert i (containsT 0 $ tree0) "one element tree should contain that elements";
-      assert i (not $ containsT 1 $ tree0) "one element tree should contain only that elements";
+      let tree0 = insertD 0 1 emptyD in
+      assert i (not $ isEmptyD tree0) "nonempty tree should not be empty";
+      assert_eq i (inr 1) (getD 0 $ tree0) "one element should be present in a one element tree";
+      assert i (containsD 0 $ tree0) "one element tree should contain that elements";
+      assert i (not $ containsD 1 $ tree0) "one element tree should contain only that elements";
     );
     group i "test {0: 1, 2: 3}" (\i.
-      let tree02 = insertT 2 3 $ insertT 0 1 emptyT in
-      assert_eq i (inr 1) (getT 0 $ tree02) "get 0 {0: 1, 2: 3} == 1";
-      assert_eq i (inr 3) (getT 2 $ tree02) "get 2 {0: 1, 2: 3} == 3";
+      let tree02 = insertD 2 3 $ insertD 0 1 emptyD in
+      assert_eq i (inr 1) (getD 0 $ tree02) "get 0 {0: 1, 2: 3} == 1";
+      assert_eq i (inr 3) (getD 2 $ tree02) "get 2 {0: 1, 2: 3} == 3";
     );
   );
 end