More map functions (#855)

stdlib/avl.mc

@@ -340,25 +340,29 @@ lang AVLTreeImpl
   sem avlMerge : all k. all a. all b. all c.
     (k -> k -> Int) -> (Option a -> Option b -> Option c) -> AVL k a ->
     AVL k b -> AVL k c
-  sem avlMerge cmp f l =
+  sem avlMerge cmp f l = | r -> avlMergeWithKey cmp (lam. f) l r
+  sem avlMergeWithKey : all k. all a. all b. all c.
+    (k -> k -> Int) -> (k -> Option a -> Option b -> Option c) -> AVL k a ->
+    AVL k b -> AVL k c
+  sem avlMergeWithKey cmp f l =
   | r ->
     match (l, r) with (Leaf _, Leaf _) then Leaf ()
     else if geqi (avlHeight l) (avlHeight r) then
       match l with Node lt then
         match avlSplit cmp lt.key r with (rl, rv, rr) in
-        let lhs = avlMerge cmp f lt.l rl in
-        let rhs = avlMerge cmp f lt.r rr in
-        match f (Some lt.value) rv with Some v then avlJoin lt.key v lhs rhs
+        let lhs = avlMergeWithKey cmp f lt.l rl in
+        let rhs = avlMergeWithKey cmp f lt.r rr in
+        match f lt.key (Some lt.value) rv with Some v then avlJoin lt.key v lhs rhs
         else avlJoin2 lhs rhs
-      else error "avlMerge: empty left tree"
+      else error "avlMergeWithKey: empty left tree"
     else
       match r with Node rt then
         match avlSplit cmp rt.key l with (ll, lv, lr) in
-        let lhs = avlMerge cmp f ll rt.l in
-        let rhs = avlMerge cmp f lr rt.r in
-        match f lv (Some rt.value) with Some v then avlJoin rt.key v lhs rhs
+        let lhs = avlMergeWithKey cmp f ll rt.l in
+        let rhs = avlMergeWithKey cmp f lr rt.r in
+        match f rt.key lv (Some rt.value) with Some v then avlJoin rt.key v lhs rhs
         else avlJoin2 lhs rhs
-      else error "avlMerge: empty right tree"
+      else error "avlMergeWithKey: empty right tree"
 
   sem avlUnionWith : all k. all v.
     (k -> k -> Int) -> (v -> v -> v) -> AVL k v -> AVL k v -> AVL k v
@@ -426,6 +430,16 @@ lang AVLTreeImpl
           avlJoin lt.key lt.value lhs rhs
       else error "avlDifference: empty left tree"
 
+  sem avlMapOption : all k. all a. all b. (k -> a -> Option b) -> AVL k a -> AVL k b
+  sem avlMapOption f =
+  | Node t ->
+    let lhs = avlMapOption f t.l in
+    let rhs = avlMapOption f t.r in
+    match f t.key t.value with Some value then avlJoin t.key value lhs rhs
+    else avlJoin2 lhs rhs
+  | Leaf _ ->
+    Leaf ()
+
   sem avlFilter : all k. all v. (k -> v -> Bool) -> AVL k v -> AVL k v
   sem avlFilter p =
   | Node t ->

stdlib/map.mc

@@ -298,12 +298,23 @@ let mapValues : all k. all v. Map k v -> [v] = lam m.
 -- │ Merge/Union │
 -- └─────────────┘
 
+-- Generalized merging of two maps. This can be used to express union,
+-- difference, intersection, etc.; any combination of two maps where
+-- we do some form of combination and filtering at each key.
 let mapMerge : all k. all a. all b. all c.
   (Option a -> Option b -> Option c) -> Map k a -> Map k b -> Map k c =
   lam f. lam l. lam r.
   use AVLTreeImpl in
   {cmp = l.cmp, root = avlMerge l.cmp f l.root r.root}
 
+-- This is `mapMerge`, except the combination function has access to
+-- the key being merged.
+let mapMergeWithKey : all k. all a. all b. all c.
+  (k -> Option a -> Option b -> Option c) -> Map k a -> Map k b -> Map k c =
+  lam f. lam l. lam r.
+  use AVLTreeImpl in
+  {cmp = l.cmp, root = avlMergeWithKey l.cmp f l.root r.root}
+
 let mapUnion : all k. all v. Map k v -> Map k v -> Map k v = lam l. lam r.
   use AVLTreeImpl in
   {l with root = avlUnionWith l.cmp (lam. lam rv. rv) l.root r.root}
@@ -332,17 +343,30 @@ let mapDifference : all k. all v. all v2. Map k v -> Map k v2 -> Map k v =
 -- │ Filter │
 -- └────────┘
 
+-- Perform a mapping and filtering at the same time, with access to
+-- the key.
+let mapMapOptionWithKey : all k. all a. all b. (k -> a -> Option b) -> Map k a -> Map k b
+  = lam f. lam m.
+  use AVLTreeImpl in
+  {root = avlMapOption f m.root, cmp = m.cmp}
+
+-- Like `mapMapOptionWithKey` but without access to the key.
+let mapMapOption : all k. all a. all b. (a -> Option b) -> Map k a -> Map k b
+  = lam f. lam m.
+  use AVLTreeImpl in
+  {root = avlMapOption (lam. f) m.root, cmp = m.cmp}
+
 -- `mapFilterWithKey p m` filters the map `m` with the predicate `p`.
 let mapFilterWithKey : all k. all v. (k -> v -> Bool) -> Map k v -> Map k v
   = lam p. lam m.
-    mapFoldWithKey
-      (lam m. lam k. lam v. if p k v then mapInsert k v m else m)
-      (mapEmpty (mapGetCmpFun m))
-      m
+  use AVLTreeImpl in
+  {root = avlFilter p m.root, cmp = m.cmp}
 
 -- `mapFilter p m` filters the map `m` with the predicate `p`.
 let mapFilter : all k. all v. (v -> Bool) -> Map k v -> Map k v
-  = lam p. mapFilterWithKey (lam. p)
+  = lam p. lam m.
+  use AVLTreeImpl in
+  {root = avlFilter (lam. p) m.root, cmp = m.cmp}
 
 mexpr
 
@@ -500,6 +524,26 @@ utest
   with [(1, "1"), (3, "3")]
 in
 
+let m = mapFromSeq subi [
+  (1, "1"),
+  (2, "2"),
+  (3, "3")
+] in
+utest
+  mapBindings (mapMapOptionWithKey (lam k. lam v. if or (eqString v "1") (eqString v "3") then Some (concat (int2string k) (cons 'x' v)) else None ()) m)
+  with [(1, "1x1"), (3, "3x3")]
+in
+
+let m = mapFromSeq subi [
+  (1, "1"),
+  (2, "2"),
+  (3, "3")
+] in
+utest
+  mapBindings (mapMapOption (lam v. if or (eqString v "1") (eqString v "3") then Some (cons 'x' v) else None ()) m)
+  with [(1, "x1"), (3, "x3")]
+in
+
 let cmp = lam a. lam b. if ltf a b then -1 else if gtf a b then 1 else 0 in
 let m = mapFromSeq cmp [(0., 0), (1., 1), (2., 2), (3., 3), (4., 4)] in
 utest mapFindUpper 4.5 m with None () in