From b57b96d13c1c62398dce29eb3b2f94335f728bd2 Mon Sep 17 00:00:00 2001
From: Lukasz Czajka <lukasz@heliax.dev>
Date: Mon, 21 Oct 2024 18:53:22 +0200
Subject: [PATCH 1/6] set functions

---
 Stdlib/Data/List.juvix      |   5 +-
 Stdlib/Data/List/Base.juvix |   2 +-
 Stdlib/Data/Map.juvix       |  45 +++++++++---
 Stdlib/Data/Set/AVL.juvix   | 143 +++++++++++++++++++++++++++---------
 4 files changed, 149 insertions(+), 46 deletions(-)

diff --git a/Stdlib/Data/List.juvix b/Stdlib/Data/List.juvix
index 6596def2..31de9aa8 100644
--- a/Stdlib/Data/List.juvix
+++ b/Stdlib/Data/List.juvix
@@ -24,7 +24,10 @@ eqListI {A} {{Eq A}} : Eq (List A) :=
       | nil nil := true
       | nil _ := false
       | _ nil := false
-      | (x :: xs) (y :: ys) := ite (Eq.eq x y) (go xs ys) false;
+      | (x :: xs) (y :: ys) :=
+        if
+          | Eq.eq x y := go xs ys
+          | else := false;
   in mkEq go;
 
 isMember {A} {{Eq A}} (elem : A) (list : List A) : Bool := isElement (Eq.eq) elem list;
diff --git a/Stdlib/Data/List/Base.juvix b/Stdlib/Data/List/Base.juvix
index e900074f..f415483c 100644
--- a/Stdlib/Data/List/Base.juvix
+++ b/Stdlib/Data/List/Base.juvix
@@ -180,7 +180,7 @@ all {A} (predicate : A -> Bool) : (list : List A) -> Bool
       | else := false;
 
 --- 𝒪(1). Returns ;true; if the ;List; is empty.
-null {A} (list : List A) : Bool :=
+isEmpty {A} (list : List A) : Bool :=
   case list of
     | nil := true
     | _ := false;
diff --git a/Stdlib/Data/Map.juvix b/Stdlib/Data/Map.juvix
index bf9083d7..9b575177 100644
--- a/Stdlib/Data/Map.juvix
+++ b/Stdlib/Data/Map.juvix
@@ -34,6 +34,7 @@ toPair {Key Value} (binding : Binding Key Value) : Pair Key Value := key binding
 instance
 bindingKeyOrdering {Key Value} {{Ord Key}} : Ord (Binding Key Value) :=
   mkOrd@{
+    {-# inline: true #-}
     cmp (b1 b2 : Binding Key Value) : Ordering := Ord.cmp (key b1) (key b2)
   };
 
@@ -58,27 +59,49 @@ insertWith
         binding := mkBinding key value;
       in mkMap (Set.insertWith fun' binding tree);
 
-insert {A B : Type} {{Ord A}} : A -> B -> Map A B -> Map A B := insertWith λ {old new := new};
+{-# inline: true #-}
+insert {Key Value : Type} {{Ord Key}} : Key -> Value -> Map Key Value -> Map Key Value :=
+  insertWith λ {old new := new};
 
 {-# specialize: [1] #-}
-lookup {A B} {{Ord A}} (k : A) : Map A B -> Maybe B
-  | (mkMap s) := map value (Set.lookupWith key k s);
+lookup {Key Value} {{Ord Key}} (lookupKey : Key) (container : Map Key Value) : Maybe Value :=
+  case container of mkMap s := map value (Set.lookupWith key lookupKey s);
 
 {-# specialize: [1, f] #-}
-fromListWith {A B} {{Ord A}} (f : B -> B -> B) (xs : List (Pair A B)) : Map A B :=
+fromListWith {Key Value} {{Ord Key}} (f : Value -> Value -> Value) (xs : List (Pair Key Value)) : Map Key Value :=
   for (acc := empty) (k, v in xs)
     insertWith f k v acc;
 
-fromList {A B} {{Ord A}} : List (Pair A B) -> Map A B := fromListWith λ {old new := new};
+{-# inline: true #-}
+fromList {Key Value} {{Ord Key}} : List (Pair Key Value) -> Map Key Value := fromListWith λ {old new := new};
 
-toList {A B} : Map A B -> List (Pair A B)
-  | (mkMap s) := map (x in Set.toList s) toPair x;
+toList {Key Value} (container : Map Key Value) : List (Pair Key Value) :=
+  case container of mkMap s := map (x in Set.toList s) toPair x;
 
-size {A B} : Map A B -> Nat
-  | (mkMap s) := Set.size s;
+{-# inline: true #-}
+size {Key Value} (container : Map Key Value) : Nat := case container of mkMap s := Set.size s;
 
-delete {A B} {{Ord A}} (k : A) : Map A B -> Map A B
-  | m@(mkMap s) := Set.deleteWith key k s |> mkMap;
+{-# inline: true #-}
+delete {Key Value} {{Ord Key}} (deleteKey : Key) (container : Map Key Value) : Map Key Value :=
+  case container of mkMap s := mkMap (Set.deleteWith key deleteKey s);
+
+keys {Key Value} (container : Map Key Value) : List Key :=
+  case container of mkMap s := map (x in Set.toList s) {key x};
+
+values {Key Value} (container : Map Key Value) : List Value :=
+  case container of mkMap s := map (x in Set.toList s) {value x};
+
+keySet {Key Value} {{Ord Key}} (container : Map Key Value) : Set Key :=
+  case container of mkMap s := Set.map (x in s) {key x};
+
+valueSet {Key Value} {{Ord Value}} (container : Map Key Value) : Set Value :=
+  case container of mkMap s := Set.map (x in s) {value x};
+
+mapValues {Key Value1 Value2} (f : Value1 -> Value2) (container : Map Key Value1) : Map Key Value2 :=
+  case container of mkMap s := mkMap (Set.Internal.unsafeMap \ {(mkBinding k v) := mkBinding k (f v)} s);
 
 instance
 eqMapI {A B} {{Eq A}} {{Eq B}} : Eq (Map A B) := mkEq (Eq.eq on toList);
+
+instance
+ordMapI {A B} {{Ord A}} {{Ord B}} : Ord (Map A B) := mkOrd (Ord.cmp on toList);
diff --git a/Stdlib/Data/Set/AVL.juvix b/Stdlib/Data/Set/AVL.juvix
index 2d6f9fc1..1be90858 100644
--- a/Stdlib/Data/Set/AVL.juvix
+++ b/Stdlib/Data/Set/AVL.juvix
@@ -14,9 +14,9 @@ import Stdlib.Data.Nat open;
 import Stdlib.Data.Int open;
 import Stdlib.Data.Bool open;
 import Stdlib.Data.Pair open;
-import Stdlib.Data.List open hiding {isMember};
+import Stdlib.Data.List open using {List; nil; ::; filter; ++; catMaybes};
 import Stdlib.Data.String open;
-import Stdlib.Trait.Foldable open;
+import Stdlib.Trait.Foldable open hiding {foldr; foldl};
 import Stdlib.Function open;
 
 --- A self-balancing binary search tree.
@@ -31,6 +31,24 @@ type AVLTree (A : Type) :=
       right : AVLTree A
     };
 
+module Internal;
+  --- 𝒪(n). Maps a function over an ;AVLTree;. Does not check if
+  --- after mapping the order of the elements is preserved. It is the
+  --- responsibility of the caller to ensure that `f` preserves the ordering.
+  unsafeMap {A B} (f : A -> B) (tree : AVLTree A) : AVLTree B :=
+    let
+      go : AVLTree A -> AVLTree B
+        | empty := empty
+        | node@{element; height; left; right} :=
+          -- TODO: we cannot use record creation syntax here because of a bug in the type-checker (https://github.com/anoma/juvix/issues/3112)
+            node
+            (f element)
+            height
+            (go left)
+            (go right);
+    in go tree;
+end;
+
 --- 𝒪(1) Retrieves the height of an ;AVLTree;. The height is the distance from
 --- the root to the deepest child.
 height {A} (tree : AVLTree A) : Nat :=
@@ -123,6 +141,7 @@ lookupWith {A K} {{order : Ord K}} (fun : A -> K) (elem : K) (tree : AVLTree A)
 lookup {A} {{Ord A}} (elem : A) (tree : AVLTree A) : Maybe A := lookupWith id elem tree;
 
 --- 𝒪(log 𝓃). Queries whether an element is in an ;AVLTree;.
+{-# specialize: [1] #-}
 isMember {A} {{Ord A}} (elem : A) (tree : AVLTree A) : Bool := isJust (lookupWith id elem tree);
 
 --- 𝒪(log 𝓃). Inserts an element elem into and ;AVLTree; using a function to
@@ -144,6 +163,7 @@ insertWith {A} {{order : Ord A}} (fun : A -> A -> A) (elem : A) (tree : AVLTree
   in go tree;
 
 --- 𝒪(log 𝓃). Inserts an element into and ;AVLTree;.
+{-# specialize: [1] #-}
 insert {A} {{Ord A}} (elem : A) (tree : AVLTree A) : AVLTree A := insertWith (flip const) elem tree;
 
 --- 𝒪(log 𝓃). Removes the minimum element from an ;AVLTree;.
@@ -229,25 +249,99 @@ size {A} : AVLTree A -> Nat
   | empty := 0
   | node@{left; right} := 1 + size left + size right;
 
--- TODO: Implement this tail-recursively with an accumulator. Then it will be O(n).
---- 𝒪(n log n). Returns the elements of an ;AVLTree; in ascending order.
-toList {A} : AVLTree A -> List A
-  | empty := nil
-  | node@{element; left; right} := toList left ++ (element :: nil) ++ toList right;
+{-# specialize: [1] #-}
+foldr {A B} (fun : A -> B -> B) (acc : B) : AVLTree A -> B
+  | empty := acc
+  | node@{element; left; right} := foldr fun (fun element (foldr fun acc right)) left;
+
+{-# specialize: [1] #-}
+foldl {A B} (fun : B -> A -> B) (acc : B) : AVLTree A -> B
+  | empty := acc
+  | node@{element; left; right} := foldl fun (fun (foldl fun acc left) element) right;
+
+{-# specialize: true, inline: case #-}
+instance
+polymorphicFoldableAVLTreeI : Polymorphic.Foldable AVLTree :=
+  Polymorphic.mkFoldable@{
+    {-# inline: true #-}
+    for {A B} (f : B -> A -> B) (acc : B) (tree : AVLTree A) : B := foldl f acc tree;
+    {-# inline: true #-}
+    rfor {A B} (f : B -> A -> B) (acc : B) (tree : AVLTree A) : B := foldr (flip f) acc tree
+  };
+
+{-# specialize: true, inline: true #-}
+instance
+foldableAVLTreeI {A} : Foldable (AVLTree A) A := fromPolymorphicFoldable;
+
+--- 𝒪(n). Returns the elements of an ;AVLTree; in ascending order.
+toList {A} (tree : AVLTree A) : List A := rfor (acc := nil) (x in tree) {x :: acc};
 
--- TODO: avoid conversion to list.
 --- 𝒪(n log n). Returns an ;AVLTree; containing elements that are members of both ;AVLTree;s.
-intersection {A} {{Ord A}} (s1 s2 : AVLTree A) : AVLTree A :=
-  toList s1 |> filter \ {x := isMember x s2} |> fromList;
+{-# specialize: [1] #-}
+intersection {A} {{Ord A}} (tree1 tree2 : AVLTree A) : AVLTree A :=
+  for (acc := empty) (x in tree1)
+    {if
+      | isMember x tree2 := insert x acc
+      | else := acc};
 
--- TODO: avoid conversion to list.
 --- 𝒪(n log n). Returns an ;AVLTree; containing elements that are members of the first but not the second ;AVLTree;.
-difference {A} {{Ord A}} (s1 s2 : AVLTree A) : AVLTree A :=
-  toList s1 |> filter \ {x := not (isMember x s2)} |> fromList;
+{-# specialize: [1] #-}
+difference {A} {{Ord A}} (tree1 tree2 : AVLTree A) : AVLTree A :=
+  for (acc := empty) (x in tree1)
+    {if
+      | isMember x tree2 := acc
+      | else := insert x acc};
 
--- TODO: avoid conversion to list.
 --- 𝒪(n log n). Returns an ;AVLTree; containing elements that are members of either the first or second ;AVLTree;.
-union {A} {{Ord A}} (s1 s2 : AVLTree A) : AVLTree A := fromList (toList s1 ++ toList s2);
+{-# specialize: [1] #-}
+union {A} {{Ord A}} (tree1 tree2 : AVLTree A) : AVLTree A :=
+  for (acc := tree1) (x in tree2) {insert x acc};
+
+syntax iterator all {init := 0; range := 1};
+--- 𝒪(𝓃). Returns ;true; if all elements of the ;AVLTree; satisfy the predicate.
+{-# specialize: [1] #-}
+all {A} (predicate : A -> Bool) (tree : AVLTree A) : Bool :=
+  let
+    go : AVLTree A -> Bool
+      | empty := true
+      | node@{element; left; right} := predicate element && go left && go right;
+  in go tree;
+
+syntax iterator any {init := 0; range := 1};
+--- 𝒪(𝓃). Returns ;true; if some elements of the ;AVLTree; satisfies the predicate.
+{-# specialize: [1] #-}
+any {A} (predicate : A -> Bool) (tree : AVLTree A) : Bool :=
+  let
+    go : AVLTree A -> Bool
+      | empty := true
+      | node@{element; left; right} := predicate element || go left || go right;
+  in go tree;
+
+syntax iterator filter {init := 0; range := 1};
+--- 𝒪(n log n). Returns an ;AVLTree; containing all elements of the tree that satisfy the predicate.
+{-# specialize: [1] #-}
+filter {A} {{Ord A}} (predicate : A -> Bool) (tree : AVLTree A) : AVLTree A :=
+  for (acc := empty) (x in tree)
+    {if
+      | predicate x := insert x acc
+      | else := acc};
+
+--- O(n log n). Partition the set into two sets, one with all elements that satisfy the predicate and one with all elements that don't satisfy the predicate.
+partition {A} {{Ord A}} (predicate : A -> Bool) (tree : AVLTree A) : Pair (AVLTree A) (AVLTree A) :=
+  for (trueSet, falseSet := empty, empty) (x in tree)
+    {if
+      | predicate x := insert x trueSet, falseSet
+      | else := trueSet, insert x falseSet};
+
+--- 𝒪(1). Creates an ;AVLTree; with a single element.
+singleton {A} (element : A) : AVLTree A := mkNode element empty empty;
+
+--- 𝒪(n log n). Checks if all elements of tree1 are in tree2.
+isSubset {A} {{Ord A}} (tree1 tree2 : AVLTree A) : Bool := all (x in tree1) {isMember x tree2};
+
+syntax iterator map {init := 0; range := 1};
+map {A B} {{Ord B}} (f : A -> B) (tree : AVLTree A) : AVLTree B :=
+  for (acc := empty) (x in tree) {insert (f x) acc};
 
 --- Formats the tree in a debug friendly format.
 prettyDebug {A} {{Show A}} (tree : AVLTree A) : String :=
@@ -275,25 +369,8 @@ toTree {A} : (tree : AVLTree A) -> Maybe (Tree A)
 --- Returns the textual representation of an ;AVLTree;.
 pretty {A} {{Show A}} (tree : AVLTree A) : String := maybe "empty" Tree.treeToString (toTree tree);
 
--- TODO: remove toList
 instance
-eqAVLTreeI {A} {{Eq A}} : Eq (AVLTree A) := mkEq ((==) on toList);
+eqAVLTreeI {A} {{Eq A}} : Eq (AVLTree A) := mkEq (Eq.eq on toList);
 
--- TODO: remove toList
 instance
 ordAVLTreeI {A} {{Ord A}} : Ord (AVLTree A) := mkOrd (Ord.cmp on toList);
-
--- TODO: remove toList
-{-# specialize: true, inline: case #-}
-instance
-polymorphicFoldableAVLTreeI : Polymorphic.Foldable AVLTree :=
-  Polymorphic.mkFoldable@{
-    for {A B} (f : B -> A -> B) (acc : B) (tree : AVLTree A) : B :=
-      Polymorphic.Foldable.for f acc (toList tree);
-    rfor {A B} (f : B → A → B) (acc : B) (tree : AVLTree A) : B :=
-      Polymorphic.Foldable.rfor f acc (toList tree)
-  };
-
-{-# specialize: true, inline: true #-}
-instance
-foldableAVLTreeI {A} : Foldable (AVLTree A) A := fromPolymorphicFoldable;

From ba7a7c0396a417819408e5258e5e11c1cd1e9f0d Mon Sep 17 00:00:00 2001
From: Lukasz Czajka <lukasz@heliax.dev>
Date: Tue, 22 Oct 2024 15:22:02 +0200
Subject: [PATCH 2/6] Map functionsd

---
 Stdlib/Data/Map.juvix     | 100 ++++++++++++++++++++----
 Stdlib/Data/Set/AVL.juvix | 155 ++++++++++++++++++++------------------
 2 files changed, 167 insertions(+), 88 deletions(-)

diff --git a/Stdlib/Data/Map.juvix b/Stdlib/Data/Map.juvix
index 9b575177..99a1b154 100644
--- a/Stdlib/Data/Map.juvix
+++ b/Stdlib/Data/Map.juvix
@@ -4,18 +4,16 @@ import Stdlib.Data.Pair open;
 import Stdlib.Data.Maybe open;
 import Stdlib.Data.List open;
 import Stdlib.Data.Nat open;
+import Stdlib.Data.Bool open;
 
 import Stdlib.Trait.Functor open;
-import Stdlib.Trait.Foldable open;
+import Stdlib.Trait.Foldable open hiding {foldr; foldl};
 import Stdlib.Trait.Ord open;
 
 import Stdlib.Function open;
 
-import Stdlib.Data.Set as Set;
-open Set using {Set};
-
-import Stdlib.Data.Set.AVL as AVL;
-open AVL using {AVLTree};
+import Stdlib.Data.Set as Set open using {Set};
+import Stdlib.Data.Set.AVL as AVL open using {AVLTree};
 
 import Stdlib.Data.BinaryTree as Tree;
 
@@ -55,32 +53,62 @@ insertWith
     mkMap tree :=
       let
         fun' (binding1 binding2 : Binding Key Value) : Binding Key Value :=
-          case binding1, binding2 of mkBinding a b1, mkBinding _ b2 := mkBinding a (fun b1 b2);
+          case binding1, binding2 of mkBinding k b1, mkBinding _ b2 := mkBinding k (fun b1 b2);
         binding := mkBinding key value;
       in mkMap (Set.insertWith fun' binding tree);
 
+--- 𝒪(log 𝓃). Inserts an element into a ;Map; at a given key.
 {-# inline: true #-}
 insert {Key Value : Type} {{Ord Key}} : Key -> Value -> Map Key Value -> Map Key Value :=
-  insertWith λ {old new := new};
+  insertWith \ {old new := new};
 
+--- 𝒪(log 𝓃). Queries whether a given key is in a ;Map;.
 {-# specialize: [1] #-}
 lookup {Key Value} {{Ord Key}} (lookupKey : Key) (container : Map Key Value) : Maybe Value :=
   case container of mkMap s := map value (Set.lookupWith key lookupKey s);
 
+--- 𝒪(log 𝓃). Queries whether a given key is in a ;Map;.
+{-# specialize: [1] #-}
+isMember {Key Value} {{Ord Key}} (key : Key) (container : Map Key Value) : Bool :=
+  isJust (lookup key container);
+
 {-# specialize: [1, f] #-}
-fromListWith {Key Value} {{Ord Key}} (f : Value -> Value -> Value) (xs : List (Pair Key Value)) : Map Key Value :=
-  for (acc := empty) (k, v in xs)
-    insertWith f k v acc;
+fromListWithKey
+  {Key Value}
+  {{Ord Key}}
+  (f : Key -> Value -> Value -> Value)
+  (xs : List (Pair Key Value))
+  : Map Key Value := for (acc := empty) (k, v in xs) {insertWith (f k) k v acc};
 
 {-# inline: true #-}
-fromList {Key Value} {{Ord Key}} : List (Pair Key Value) -> Map Key Value := fromListWith λ {old new := new};
+fromListWith
+  {Key Value}
+  {{Ord Key}}
+  (f : Value -> Value -> Value)
+  (xs : List (Pair Key Value))
+  : Map Key Value := fromListWithKey (const f) xs;
+
+{-# inline: true #-}
+fromList {Key Value} {{Ord Key}} : List (Pair Key Value) -> Map Key Value :=
+  fromListWith λ {old new := new};
 
 toList {Key Value} (container : Map Key Value) : List (Pair Key Value) :=
-  case container of mkMap s := map (x in Set.toList s) toPair x;
+  case container of mkMap s := map (x in Set.toList s) {toPair x};
+
+{-# specialize: [1, f] #-}
+fromSet {Key Value} {{Ord Key}} (f : Key -> Value) (set : Set Key) : Map Key Value :=
+  for (acc := empty) (k in set) {insert k (f k) acc};
 
+--- 𝒪(1). Checks if a ;Map; is empty.
+{-# inline: true #-}
+isEmpty {Key Value} (container : Map Key Value) : Bool :=
+  case container of mkMap s := Set.isEmpty s;
+
+--- 𝒪(𝓃). Returns the number of elements of a ;Map;.
 {-# inline: true #-}
 size {Key Value} (container : Map Key Value) : Nat := case container of mkMap s := Set.size s;
 
+--- 𝒪(log 𝓃). Removes a key assignment from a ;Map;.
 {-# inline: true #-}
 delete {Key Value} {{Ord Key}} (deleteKey : Key) (container : Map Key Value) : Map Key Value :=
   case container of mkMap s := mkMap (Set.deleteWith key deleteKey s);
@@ -97,8 +125,50 @@ keySet {Key Value} {{Ord Key}} (container : Map Key Value) : Set Key :=
 valueSet {Key Value} {{Ord Value}} (container : Map Key Value) : Set Value :=
   case container of mkMap s := Set.map (x in s) {value x};
 
-mapValues {Key Value1 Value2} (f : Value1 -> Value2) (container : Map Key Value1) : Map Key Value2 :=
-  case container of mkMap s := mkMap (Set.Internal.unsafeMap \ {(mkBinding k v) := mkBinding k (f v)} s);
+{-# specialize: [f] #-}
+mapValuesWithKey
+  {Key Value1 Value2} (f : Key -> Value1 -> Value2) (container : Map Key Value1) : Map Key Value2 :=
+  case container of
+    mkMap s := mkMap (Set.Internal.unsafeMap \ {(mkBinding k v) := mkBinding k (f k v)} s);
+
+{-# inline: true #-}
+mapValues
+  {Key Value1 Value2} (f : Value1 -> Value2) (container : Map Key Value1) : Map Key Value2 :=
+  mapValuesWithKey (const f) container;
+
+{-# inline: true #-}
+singleton {Key Value} {{Ord Key}} (key : Key) (value : Value) : Map Key Value :=
+  insert key value empty;
+
+{-# inline: true #-}
+foldr
+  {Key Value Acc} (f : Key -> Value -> Acc -> Acc) (acc : Acc) (container : Map Key Value) : Acc :=
+  case container of mkMap s := rfor (acc := acc) (x in s) {f (key x) (value x) acc};
+
+{-# inline: true #-}
+foldl
+  {Key Value Acc} (f : Acc -> Key -> Value -> Acc) (acc : Acc) (container : Map Key Value) : Acc :=
+  case container of mkMap s := for (acc := acc) (x in s) {f acc (key x) (value x)};
+
+{-# specialize: true, inline: case #-}
+instance
+functorMapI {Key} : Functor (Map Key) :=
+  mkFunctor@{
+    map {A B} (f : A -> B) (container : Map Key A) : Map Key B := mapValues f container
+  };
+
+{-# specialize: true, inline: case #-}
+instance
+foldableMapI {Key Value} : Foldable (Map Key Value) (Pair Key Value) :=
+  mkFoldable@{
+    {-# inline: true #-}
+    for {B} (f : B -> Pair Key Value -> B) (acc : B) (container : Map Key Value) : B :=
+      foldl \ {a k v := f a (k, v)} acc container;
+
+    {-# inline: true #-}
+    rfor {B} (f : B -> Pair Key Value -> B) (acc : B) (container : Map Key Value) : B :=
+      foldr \ {k v a := f a (k, v)} acc container
+  };
 
 instance
 eqMapI {A B} {{Eq A}} {{Eq B}} : Eq (Map A B) := mkEq (Eq.eq on toList);
diff --git a/Stdlib/Data/Set/AVL.juvix b/Stdlib/Data/Set/AVL.juvix
index 1be90858..ae82ca97 100644
--- a/Stdlib/Data/Set/AVL.juvix
+++ b/Stdlib/Data/Set/AVL.juvix
@@ -31,6 +31,13 @@ type AVLTree (A : Type) :=
       right : AVLTree A
     };
 
+--- 𝒪(1) Retrieves the height of an ;AVLTree;. The height is the distance from
+--- the root to the deepest child.
+height {A} (tree : AVLTree A) : Nat :=
+  case tree of
+    | empty := 0
+    | node@{height} := height;
+
 module Internal;
   --- 𝒪(n). Maps a function over an ;AVLTree;. Does not check if
   --- after mapping the order of the elements is preserved. It is the
@@ -47,79 +54,74 @@ module Internal;
             (go left)
             (go right);
     in go tree;
-end;
-
---- 𝒪(1) Retrieves the height of an ;AVLTree;. The height is the distance from
---- the root to the deepest child.
-height {A} (tree : AVLTree A) : Nat :=
-  case tree of
-    | empty := 0
-    | node@{height} := height;
-
-type BalanceFactor :=
-  | --- Left child is higher.
-    LeanLeft
-  | --- Equal heights of children.
-    LeanNone
-  | --- Right child is higher.
-    LeanRight;
 
-diffFactor {A} (tree : AVLTree A) : Int :=
-  case tree of
-    | empty := 0
-    | node@{left; right} := intSubNat (height right) (height left);
-
-{-# inline: true #-}
-balanceFactor' {A} (left right : AVLTree A) : BalanceFactor :=
-  let
-    h1 := height left;
-    h2 := height right;
-  in if
-    | h1 < h2 := LeanRight
-    | h2 < h1 := LeanLeft
-    | else := LeanNone;
-
---- 𝒪(1). Computes the balance factor, i.e., the height of the right subtree
---- minus the height of the left subtree.
-balanceFactor {A} (tree : AVLTree A) : BalanceFactor :=
-  -- We avoid `diffFactor` here for efficiency.
-  case tree of
-    | empty := LeanNone
-    | node@{left; right} := balanceFactor' left right;
-
---- 𝒪(1). Helper function for creating a node.
-mkNode {A} (value : A) (left : AVLTree A) (right : AVLTree A) : AVLTree A :=
-  node value (1 + max (height left) (height right)) left right;
-
---- 𝒪(1). Left rotation.
-rotateLeft {A} (tree : AVLTree A) : AVLTree A :=
-  case tree of
-    | node x _ a (node z _ b c) := mkNode z (mkNode x a b) c
-    | n := n;
-
---- 𝒪(1). Right rotation.
-rotateRight {A} (tree : AVLTree A) : AVLTree A :=
-  case tree of
-    | node z _ (node y _ x t3) t4 := mkNode y x (mkNode z t3 t4)
-    | n := n;
+  type BalanceFactor :=
+    | --- Left child is higher.
+      LeanLeft
+    | --- Equal heights of children.
+      LeanNone
+    | --- Right child is higher.
+      LeanRight;
+
+  diffFactor {A} (tree : AVLTree A) : Int :=
+    case tree of
+      | empty := 0
+      | node@{left; right} := intSubNat (height right) (height left);
+
+  {-# inline: true #-}
+  balanceFactor' {A} (left right : AVLTree A) : BalanceFactor :=
+    let
+      h1 := height left;
+      h2 := height right;
+    in if
+      | h1 < h2 := LeanRight
+      | h2 < h1 := LeanLeft
+      | else := LeanNone;
+
+  --- 𝒪(1). Computes the balance factor, i.e., the height of the right subtree
+  --- minus the height of the left subtree.
+  balanceFactor {A} (tree : AVLTree A) : BalanceFactor :=
+    -- We avoid `diffFactor` here for efficiency.
+    case tree of
+      | empty := LeanNone
+      | node@{left; right} := balanceFactor' left right;
+
+  --- 𝒪(1). Helper function for creating a node.
+  mkNode {A} (value : A) (left : AVLTree A) (right : AVLTree A) : AVLTree A :=
+    node value (1 + max (height left) (height right)) left right;
+
+  --- 𝒪(1). Left rotation.
+  rotateLeft {A} (tree : AVLTree A) : AVLTree A :=
+    case tree of
+      | node x _ a (node z _ b c) := mkNode z (mkNode x a b) c
+      | n := n;
+
+  --- 𝒪(1). Right rotation.
+  rotateRight {A} (tree : AVLTree A) : AVLTree A :=
+    case tree of
+      | node z _ (node y _ x t3) t4 := mkNode y x (mkNode z t3 t4)
+      | n := n;
+
+  --- 𝒪(1). Applies local rotations if needed.
+  balance {A} (tree : AVLTree A) : AVLTree A :=
+    case tree of
+      | empty := empty
+      | node@{element; height; left; right} :=
+        case balanceFactor' left right of
+          | LeanRight :=
+            case balanceFactor right of {
+              | LeanLeft := rotateLeft (mkNode element left (rotateRight right))
+              | _ := rotateLeft tree
+            }
+          | LeanLeft :=
+            case balanceFactor left of {
+              | LeanRight := rotateRight (mkNode element (rotateLeft left) right)
+              | _ := rotateRight tree
+            }
+          | LeanNone := tree;
+end;
 
---- 𝒪(1). Applies local rotations if needed.
-balance {A} (tree : AVLTree A) : AVLTree A :=
-  case tree of
-    | empty := empty
-    | node@{element; height; left; right} :=
-      case balanceFactor' left right of
-        | LeanRight :=
-          case balanceFactor right of {
-            | LeanLeft := rotateLeft (mkNode element left (rotateRight right))
-            | _ := rotateLeft tree
-          }
-        | LeanLeft :=
-          case balanceFactor left of {
-            | LeanRight := rotateRight (mkNode element (rotateLeft left) right)
-            | _ := rotateRight tree
-          }
-        | LeanNone := tree;
+open Internal;
 
 --- 𝒪(log 𝓃). Lookup a member from the ;AVLTree; using a projection function.
 --- Ord A, Ord K and fun must be compatible. i.e cmp_k (fun a1) (fun a2) == cmp_a a1 a2
@@ -162,7 +164,7 @@ insertWith {A} {{order : Ord A}} (fun : A -> A -> A) (elem : A) (tree : AVLTree
           | Equal := node (fun element elem) height left right;
   in go tree;
 
---- 𝒪(log 𝓃). Inserts an element into and ;AVLTree;.
+--- 𝒪(log 𝓃). Inserts an element into an ;AVLTree;.
 {-# specialize: [1] #-}
 insert {A} {{Ord A}} (elem : A) (tree : AVLTree A) : AVLTree A := insertWith (flip const) elem tree;
 
@@ -206,7 +208,7 @@ deleteWith
 
 --- 𝒪(log 𝓃). Removes an element from an ;AVLTree;.
 {-# specialize: [1] #-}
-delete {A} {{o : Ord A}} : A -> AVLTree A -> AVLTree A := deleteWith id;
+delete {A} {{Ord A}} : A -> AVLTree A -> AVLTree A := deleteWith id;
 
 --- 𝒪(log 𝓃). Returns the minimum element of the ;AVLTree;.
 lookupMin {A} : AVLTree A -> Maybe A
@@ -244,6 +246,13 @@ isBalanced {A} : (tree : AVLTree A) -> Bool
 {-# specialize: [1] #-}
 fromList {A} {{Ord A}} (list : List A) : AVLTree A := for (acc := empty) (x in list) {insert x acc};
 
+--- 𝒪(1). Checks if an ;AVLTree; is empty.
+{-# inline: true #-}
+isEmpty {A} (tree : AVLTree A) : Bool :=
+  case tree of
+    | empty := true
+    | node@{} := false;
+
 --- 𝒪(𝓃). Returns the number of elements of an ;AVLTree;.
 size {A} : AVLTree A -> Nat
   | empty := 0

From 4de9c4665da912cd44d0ffd6cc3b326c03a79380 Mon Sep 17 00:00:00 2001
From: Lukasz Czajka <lukasz@heliax.dev>
Date: Tue, 22 Oct 2024 15:26:51 +0200
Subject: [PATCH 3/6] fix tests

---
 test/Test.juvix     | 8 ++++----
 test/Test/AVL.juvix | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/test/Test.juvix b/test/Test.juvix
index f52e4311..1409f631 100644
--- a/test/Test.juvix
+++ b/test/Test.juvix
@@ -22,7 +22,7 @@ propReverseDoesNotChangeLength (xs : List Int) : Bool := length (reverse xs) ==
 propReverseReverseIsIdentity (xs : List Int) : Bool := eqListInt xs (reverse (reverse xs));
 
 propTailLengthOneLess (xs : List Int) : Bool :=
-  null xs || ofNat (length (tail xs)) == intSubNat (length xs) 1;
+  isEmpty xs || ofNat (length (tail xs)) == intSubNat (length xs) 1;
 
 propSplitAtRecombine (n : Nat) (xs : List Int) : Bool :=
   case splitAt n xs of lhs, rhs := eqListInt xs (lhs ++ rhs);
@@ -95,11 +95,11 @@ propTransposeMatrixDimensions (xs : List (List Int)) : Bool :=
     checkTxsRowXsCol : Bool :=
       case xs of
         | x :: _ := length txs == length x
-        | _ := null txs;
+        | _ := isEmpty txs;
     checkXsRowTxsCol : Bool :=
       case txs of
         | tx :: _ := length xs == length tx
-        | _ := null xs;
+        | _ := isEmpty xs;
   in checkTxsRowXsCol && checkXsRowTxsCol;
 
 propFoundElementSatisfiesPredicate (p : Int -> Bool) (xs : List Int) : Bool :=
@@ -129,7 +129,7 @@ propFindWithEmptyList (p : Int -> Bool) : Bool := find p [] == nothing;
 
 propFindWithAlwaysTrueIsJust (xs : List Int) : Bool :=
   if
-    | null xs := true
+    | isEmpty xs := true
     | else :=
       case find (const true) xs of
         | just _ := true
diff --git a/test/Test/AVL.juvix b/test/Test/AVL.juvix
index ba9845c7..5e5eacee 100644
--- a/test/Test/AVL.juvix
+++ b/test/Test/AVL.juvix
@@ -49,7 +49,7 @@ s2DeleteWith : TestDescr Box :=
   mkTestDescr@{
     testTitle := TestDescr.testTitle s2 ++str "-delete-with";
     testLen := sub (TestDescr.testLen s2) 2;
-    testSet := deleteManyWith Box.b [1; 8] (TestDescr.testSet s2 |> toList |> map mkBox |> fromList)
+    testSet := deleteManyWith Box.b [1; 8] (TestDescr.testSet s2 |> AVL.map mkBox)
   };
 
 s3 : TestDescr Nat :=

From 8c8e91e8bf98a1f7ee19ce09ed535c6705163a4a Mon Sep 17 00:00:00 2001
From: Lukasz Czajka <lukasz@heliax.dev>
Date: Tue, 22 Oct 2024 15:43:49 +0200
Subject: [PATCH 4/6] Map functions

---
 Stdlib/Data/Map.juvix | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/Stdlib/Data/Map.juvix b/Stdlib/Data/Map.juvix
index 99a1b154..b9675316 100644
--- a/Stdlib/Data/Map.juvix
+++ b/Stdlib/Data/Map.juvix
@@ -150,6 +150,30 @@ foldl
   {Key Value Acc} (f : Acc -> Key -> Value -> Acc) (acc : Acc) (container : Map Key Value) : Acc :=
   case container of mkMap s := for (acc := acc) (x in s) {f acc (key x) (value x)};
 
+syntax iterator all {init := 0; range := 1};
+--- 𝒪(𝓃). Returns ;true; if all key-element pairs in the ;Map; satisfy the predicate.
+{-# inline: true #-}
+all {Key Value} (predicate : Key -> Value -> Bool) (container : Map Key Value) : Bool :=
+  case container of mkMap s := Set.all (x in s) {predicate (key x) (value x)};
+
+syntax iterator any {init := 0; range := 1};
+--- 𝒪(𝓃). Returns ;true; if some key-element pair in the ;Map; satisfies the predicate.
+{-# inline: true #-}
+any {Key Value} (predicate : Key -> Value -> Bool) (container : Map Key Value) : Bool :=
+  case container of mkMap s := Set.any (x in s) {predicate (key x) (value x)};
+
+syntax iterator filter {init := 0; range := 1};
+--- 𝒪(n log n). Returns a ;Map; containing all key-element pairs of the given
+--- map container that satisfy the predicate.
+{-# inline: true #-}
+filter
+  {Key Value}
+  {{Ord Key}}
+  (predicate : Key -> Value -> Bool)
+  (container : Map Key Value)
+  : Map Key Value :=
+  case container of mkMap s := mkMap (Set.filter (x in s) {predicate (key x) (value x)});
+
 {-# specialize: true, inline: case #-}
 instance
 functorMapI {Key} : Functor (Map Key) :=

From 72caf3dec42337285c37851cad5b2e8f7fd2fd44 Mon Sep 17 00:00:00 2001
From: Lukasz Czajka <lukasz@heliax.dev>
Date: Tue, 22 Oct 2024 17:47:04 +0200
Subject: [PATCH 5/6] tests

---
 Stdlib/Data/Map.juvix     | 13 ++++++++++
 Stdlib/Data/Set/AVL.juvix |  1 +
 test/Test/AVL.juvix       | 27 ++++++++++++++++++++-
 test/Test/Map.juvix       | 50 +++++++++++++++++++++++++++++++++------
 4 files changed, 83 insertions(+), 8 deletions(-)

diff --git a/Stdlib/Data/Map.juvix b/Stdlib/Data/Map.juvix
index b9675316..d441d85e 100644
--- a/Stdlib/Data/Map.juvix
+++ b/Stdlib/Data/Map.juvix
@@ -174,6 +174,19 @@ filter
   : Map Key Value :=
   case container of mkMap s := mkMap (Set.filter (x in s) {predicate (key x) (value x)});
 
+syntax iterator partition {init := 0; range := 1};
+{-# inline: true #-}
+partition
+  {Key Value}
+  {{Ord Key}}
+  (predicate : Key -> Value -> Bool)
+  (container : Map Key Value)
+  : Pair (Map Key Value) (Map Key Value) :=
+  case container of
+    mkMap s :=
+      case Set.partition (x in s) {predicate (key x) (value x)} of
+        left, right := mkMap left, mkMap right;
+
 {-# specialize: true, inline: case #-}
 instance
 functorMapI {Key} : Functor (Map Key) :=
diff --git a/Stdlib/Data/Set/AVL.juvix b/Stdlib/Data/Set/AVL.juvix
index ae82ca97..2c0619b9 100644
--- a/Stdlib/Data/Set/AVL.juvix
+++ b/Stdlib/Data/Set/AVL.juvix
@@ -336,6 +336,7 @@ filter {A} {{Ord A}} (predicate : A -> Bool) (tree : AVLTree A) : AVLTree A :=
       | else := acc};
 
 --- O(n log n). Partition the set into two sets, one with all elements that satisfy the predicate and one with all elements that don't satisfy the predicate.
+syntax iterator partition {init := 0; range := 1};
 partition {A} {{Ord A}} (predicate : A -> Bool) (tree : AVLTree A) : Pair (AVLTree A) (AVLTree A) :=
   for (trueSet, falseSet := empty, empty) (x in tree)
     {if
diff --git a/test/Test/AVL.juvix b/test/Test/AVL.juvix
index 5e5eacee..33670451 100644
--- a/test/Test/AVL.juvix
+++ b/test/Test/AVL.juvix
@@ -2,6 +2,7 @@ module Test.AVL;
 
 import Test.JuvixUnit open;
 import Stdlib.Prelude open;
+import Stdlib.Debug as Debug;
 
 import Stdlib.Data.Set.AVL as AVL open;
 
@@ -67,7 +68,7 @@ s4 : TestDescr Nat :=
   };
 
 tests : List Test :=
-  [ testCase "s1-member" (assertTrue "member? 3 s1" (AVL.isMember 3 s1Tree))
+  [ testCase "s1-member" (assertTrue "isMember 3 s1" (AVL.isMember 3 s1Tree))
   ; testCase
     "s1-s2-intersection"
     (assertEqual "intersection s1Tree s2Tree" (intersection s1Tree s2Tree) (fromList [1; 2; 3; 8]))
@@ -77,6 +78,30 @@ tests : List Test :=
   ; testCase
     "s1-s2-union"
     (assertEqual "union s1Tree s2Tree" (union s1Tree s2Tree) (fromList [0; 1; 2; 3; 4; 8; 9]))
+  ; testCase
+    "s2-filter"
+    (assertEqual "filter (> 3) s2" [0; 1; 2] (AVL.filter ((>) 3) s2Tree |> AVL.toList))
+  ; testCase "s2-all" (assertTrue "all (> 9) s2" (AVL.all ((>) 9) s2Tree))
+  ; testCase "s2-any" (assertTrue "any (< 3) s2" (AVL.any ((<) 3) s2Tree))
+  ; testCase
+    "s2-partition"
+    (assertEqual
+      "partition (< 3) s2"
+      ([0; 1; 2], [3; 4; 8])
+      (AVL.partition ((>) 3) s2Tree |> both AVL.toList))
+  ; testCase "s2-lookup" (assertJust "lookup 4 s2" (lookup 4 s2Tree))
+  ; testCase "s2-lookup" (assertNothing (const "lookup 5 s2") (lookup 5 s2Tree))
+  ; testCase "s2-lookup" (assertNothing (const "lookup 6 s2") (lookup 6 s2Tree))
+  ; testCase "s1-s2-not-isSubset" (assertFalse "isSubset s1 s2" (isSubset s1Tree s2Tree))
+  ; testCase "s2-s1-not-isSubset" (assertFalse "isSubset s2 s1" (isSubset s2Tree s1Tree))
+  ; testCase
+    "singleton-s2-isSubset"
+    (assertTrue "isSubset (singleton 1) s2" (isSubset (singleton 1) s2Tree))
+  ; testCase "s2-foldr" (assertEqual "foldr (+) 0 s2" 18 (AVL.foldr (+) 0 s2Tree))
+  ; testCase "s2-foldl" (assertEqual "foldl (+) 0 s2" 18 (AVL.foldl (+) 0 s2Tree))
+  ; testCase
+    "s2-map"
+    (assertEqual "map (+ 1) s2" [1; 2; 3; 4; 5; 9] (AVL.map ((+) 1) s2Tree |> AVL.toList))
   ; testCase "for ascending iteration"
     <| assertEqual "for iterates in ascending order" [1; 2; 3; 4]
     <| for (acc := []) (k in fromList [3; 2; 4; 1]) {snoc acc k}
diff --git a/test/Test/Map.juvix b/test/Test/Map.juvix
index 1cf9c622..7865afdc 100644
--- a/test/Test/Map.juvix
+++ b/test/Test/Map.juvix
@@ -1,14 +1,10 @@
 module Test.Map;
 
-import Juvix.Builtin.V1 open;
-import Stdlib.System.IO open;
-import Stdlib.Data.Pair open;
-import Stdlib.Data.List open;
-import Stdlib.Function open;
+import Stdlib.Prelude open;
 import Test.JuvixUnit open;
 
-import Stdlib.Data.Map as Map;
-open Map using {Map};
+import Stdlib.Data.Map as Map open using {Map};
+import Stdlib.Data.Set as Set open using {Set};
 
 tests : List Test :=
   let
@@ -49,6 +45,46 @@ tests : List Test :=
      ; testCase
        "Map.fromList distinct keys"
        (assertEqListPair (Map.toList (Map.fromList [1, 1; 2, 2])) [1, 1; 2, 2])
+     ; testCase
+       "Map.fromSet"
+       (assertEqListPair (Map.toList (Map.fromSet id (Set.fromList [1; 2; 3]))) [1, 1; 2, 2; 3, 3])
+     ; testCase "Map.keys" (assertEqual "expected keys" (Map.keys m) [1; 2])
+     ; testCase "Map.values" (assertEqual "expected values" (Map.values m) ["one"; "two"])
+     ; testCase "Map.keySet" (assertEqual "expected key set" (Map.keySet m) (Set.fromList [1; 2]))
+     ; testCase
+       "Map.valueSet"
+       (assertEqual "expected value set" (Map.valueSet m2) (Set.fromList [2; 4]))
+     ; testCase
+       "Map.isEmpty"
+       (assertEqual "expected empty map" (Map.isEmpty (Map.empty {Nat} {Nat})) true)
+     ; testCase "Map.isEmpty" (assertEqual "expected non-empty map" (Map.isEmpty m) false)
+     ; testCase "Map.isMember" (assertEqual "expected member" (Map.isMember 1 m) true)
+     ; testCase "Map.isMember" (assertEqual "expected non-member" (Map.isMember 3 m) false)
+     ; testCase
+       "Map.filter"
+       (assertEqListPair (Map.toList (Map.filter \ {k v := k + v > 3} m2)) [3, 4])
+     ; testCase
+       "Map.partition"
+       (assertEqual
+         "expected partition"
+         ([1, 2], [3, 4])
+         (Map.partition \ {k v := k < 3} m2 |> both Map.toList))
+     ; testCase
+       "Map.mapValuesWithKey"
+       (assertEqListPair (Map.toList (Map.mapValuesWithKey \ {k v := k + v} m2)) [1, 3; 3, 7])
+     ; testCase
+       "Map.mapValues"
+       (assertEqListPair (Map.toList (Map.mapValues ((+) 1) m2)) [1, 3; 3, 5])
+     ; testCase "Map.map" (assertEqListPair (Map.toList (map ((+) 1) m2)) [1, 3; 3, 5])
+     ; testCase "Map.all" (assertTrue "expected all to be true" (Map.all \ {_ v := v > 0} m2))
+     ; testCase "Map.any" (assertTrue "expected any to be true" (Map.any \ {_ v := v > 3} m2))
+     ; testCase
+       "Map.foldl"
+       (assertEqual "expected foldl" (Map.foldl \ {acc k v := acc + k + v} 0 m2) 10)
+     ; testCase
+       "Map.foldr"
+       (assertEqual "expected foldr" (Map.foldr \ {acc k v := acc + k + v} 0 m2) 10)
+     ; testCase "Map.singleton" (assertEqListPair (Map.toList (Map.singleton 1 2)) [1, 2])
      ];
 
 suite : TestSuite := testSuite "Map" tests;

From 194e8a2d236d80f8bc2032f5ea472bef15846ec9 Mon Sep 17 00:00:00 2001
From: Lukasz Czajka <lukasz@heliax.dev>
Date: Tue, 22 Oct 2024 18:23:34 +0200
Subject: [PATCH 6/6] fix formatting

---
 Stdlib/Data/Set/AVL.juvix | 1 -
 1 file changed, 1 deletion(-)

diff --git a/Stdlib/Data/Set/AVL.juvix b/Stdlib/Data/Set/AVL.juvix
index 2c0619b9..9e837825 100644
--- a/Stdlib/Data/Set/AVL.juvix
+++ b/Stdlib/Data/Set/AVL.juvix
@@ -335,7 +335,6 @@ filter {A} {{Ord A}} (predicate : A -> Bool) (tree : AVLTree A) : AVLTree A :=
       | predicate x := insert x acc
       | else := acc};
 
---- O(n log n). Partition the set into two sets, one with all elements that satisfy the predicate and one with all elements that don't satisfy the predicate.
 syntax iterator partition {init := 0; range := 1};
 partition {A} {{Ord A}} (predicate : A -> Bool) (tree : AVLTree A) : Pair (AVLTree A) (AVLTree A) :=
   for (trueSet, falseSet := empty, empty) (x in tree)