Add intersectWithKey for map and avl

src/stdlib/avl.mc

@@ -390,29 +390,34 @@ lang AVLTreeImpl
   sem avlIntersectWith : all k. all a. all b. all c.
     (k -> k -> Int) -> (a -> b -> c) -> AVL k a -> AVL k b -> AVL k c
   sem avlIntersectWith cmp f l =
+  | r ->
+    avlIntersectWithKey cmp (lam k. f) l r
+  sem avlIntersectWithKey : all k. all a. all b. all c.
+    (k -> k -> Int) -> (k -> a -> b -> c) -> AVL k a -> AVL k b -> AVL k c
+  sem avlIntersectWithKey cmp f l =
   | r ->
     match l with Leaf _ then Leaf ()
     else match r with 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 = avlIntersectWith cmp f lt.l rl in
-        let rhs = avlIntersectWith cmp f lt.r rr in
+        let lhs = avlIntersectWithKey cmp f lt.l rl in
+        let rhs = avlIntersectWithKey cmp f lt.r rr in
         match rv with Some x then
-          avlJoin lt.key (f lt.value x) lhs rhs
+          avlJoin lt.key (f lt.key lt.value x) lhs rhs
         else
           avlJoin2 lhs rhs
-      else error "avlIntersectWith: empty left tree"
+      else error "avlIntersectWithKey: empty left tree"
     else
       match r with Node rt then
         match avlSplit cmp rt.key l with (ll, lv, lr) in
-        let lhs = avlIntersectWith cmp f ll rt.l in
-        let rhs = avlIntersectWith cmp f lr rt.r in
+        let lhs = avlIntersectWithKey cmp f ll rt.l in
+        let rhs = avlIntersectWithKey cmp f lr rt.r in
         match lv with Some x then
-          avlJoin rt.key (f x rt.value) lhs rhs
+          avlJoin rt.key (f rt.key x rt.value) lhs rhs
         else
           avlJoin2 lhs rhs
-      else error "avlIntersectWith: empty right tree"
+      else error "avlIntersectWithKey: empty right tree"
 
   sem avlDifference : all k. all a. all b. (k -> k -> Int) -> AVL k a -> AVL k b -> AVL k a
   sem avlDifference cmp l =

src/stdlib/map.mc

@@ -333,6 +333,11 @@ let mapIntersectWith : all k. all a. all b. all c. (a -> b -> c) -> Map k a -> M
   use AVLTreeImpl in
   {cmp = l.cmp, root = avlIntersectWith l.cmp f l.root r.root}
 
+let mapIntersectWithKey : all k. all a. all b. all c. (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c =
+  lam f. lam l. lam r.
+  use AVLTreeImpl in
+  {cmp = l.cmp, root = avlIntersectWithKey l.cmp f l.root r.root}
+
 let mapDifference : all k. all v. all v2. Map k v -> Map k v2 -> Map k v =
   lam l. lam r.
   use AVLTreeImpl in