funclang series4: Implemented remove and removeAll functions for ass9.

funclang-taddeus/series4/ass9.ml

 open List
 
-(* A Node has a key, value and list of children *)
+(* A trie node has a key, a value and a list of children *)
 type ('a, 'b) trie = Empty | Node of 'a list * 'b * ('a, 'b) trie list
 
-let rec starts_with l s =
-    (length s) <= (length l) && match s with
+(* Helper function, checks if a list l start with a sublist s *)
+let rec starts_with l s = match s with
     | [] -> true
-    | h::t -> h = (hd l) && starts_with (tl l) t
-;;
+    | hs::ts -> match l with
+        | [] -> false
+        | hl::tl -> hs = hl && starts_with tl ts
 
 (* Insert a new (key, value) pair into a trie *)
-let rec insert trie key value =
-    match trie with
+(* TODO: this doesn't work completely yet *)
+let rec insert trie key value = match trie with
     | Empty -> Node (key, value, []) (* Root *)
     | Node (k, v, children) ->
         if k = key then
             raise (Failure "Inserted key already exists in trie")
-        else if starts_with key k then
+        else if not (starts_with key k) then
+            raise (Failure "Inserted key does not start with node key")
+        else
             (* Inserted key should be in this node because it starts with
              * the node's key *)
             let rec walk_nodes = function
                 | [] -> Empty
-                | node::tail ->
-                    match node with
+                | node::tail -> match node with
                     | Empty -> Empty
                     | Node (k, _, _) ->
                         if starts_with key k then
@@ -45,13 +47,52 @@ let rec insert trie key value =
                 let move, siblings = partition matches_key children in
                 Node (k, v, (Node (key, value, move))::siblings)
             | node -> node
+
+(* Helper function, checks if a trie is not empty *)
+let non_empty = function | Empty -> false | _ -> true
+
+(* Remove a key from a given trie *)
+let rec remove trie key = match trie with
+    | Empty -> Empty
+    | Node (k, v, children) ->
+        if k = key then
+            (* Key match, remove the node *)
+            Empty
+        else if starts_with key k then
+            (* No full key match but prefix does match, match all children *)
+            let match_with_key node = match node with
+                | Empty -> Empty
+                | Node (k, _, _) -> if k = key then Empty else remove node key
+            in
+            Node (k, v, (filter non_empty (map match_with_key children)))
         else
-            raise (Failure "Inserted key does not start with node key")
-;;
+            (* No prefix match, so no need to check child nodes *)
+            trie
+
+(* Remove all keys that start with a given prefix from a trie.
+ * Basically does the same as the 'remove' function above, only using
+ * 'starts_with' instead of literally matching the key *)
+let rec removeAll trie prefix = match trie with
+    | Empty -> Empty
+    | Node (k, v, children) ->
+        if starts_with k prefix then
+            (* Prefix match, remove the node *)
+            Empty
+        else if starts_with prefix k then
+            (* No full prefix match yet, but possibly in children so continue
+             * matching them *)
+            let match_with_prefix node = match node with
+                | Empty -> Empty
+                | Node (key, _, _) -> if starts_with key prefix then Empty
+                                      else removeAll node prefix
+            in
+            Node (k, v, (filter non_empty (map match_with_prefix children)))
+        else
+            (* Prefix is not going to match in this subtree, don't remove it *)
+            trie
 
 (* Look up a value associated with a key *)
-let rec lookup trie key =
-    match trie with
+let rec lookup trie key = match trie with
     | Empty -> None (* Trie is empty, so no result *)
     | Node (k, value, children) ->
         if k = key then
@@ -70,11 +111,9 @@ let rec lookup trie key =
         else
             (* Not found *)
             None
-;;
 
 (* Find all (key, value) pairs whose key start with a given prefix *)
-let rec matches trie key =
-    match trie with
+let rec matches trie key = match trie with
     | Empty -> []
     | Node (k, value, children) ->
         let rec match_nodes = function
@@ -83,8 +122,3 @@ let rec matches trie key =
         in
         (if starts_with k key then [(k, value)] else [])
             @ match_nodes children
-;;
-
-(*let a = Node ([1;2;3], "a", []);;
-let b = Node ([1;3], "b", [a]);;
-let c = Node ([3;4], "c", [b]);;*)