Insert function for a trie in Ocaml

Question

I don't know how could I write smarter/more elegant code for the add function , it just seems so meaty and hard to read. This trie only stores words and marks their end with a bool.

type trie = Node of bool * (char * trie) list 
let empty = Node (false, [])
let rec add w tr = match w, tr with 
 | [], Node (_, lvl_l) -> Node (true, lvl_l)
 | wh :: wt, Node (b, lvl_l) -> 
 try let ins_point = (List.assoc wh lvl_l) in 
 Node(b, (wh, add wt ins_point)::(List.remove_assoc wh lvl_l)) 
 with Not_found -> Node (b, (wh, add wt empty)::lvl_l)
let explode word = 
 let rec explode' i acc = 
 if i < 0 then acc else explode' (i-1) (word.[i] :: acc) 
 in explode' (String.length word - 1) []

Output Example :

# let x = add (explode "horror") empty;;
val x : trie =
 Node (false,
 [('h',
 Node (false,
 [('o',
 Node (false,
 [('r',
 Node (false,
 [('r',
 Node (false, [('o', Node (false, [('r', Node (true, []))]))]))]))]))]))])
# let x = add (explode "horrific") x;;
val x : trie =
 Node (false,
 [('h',
 Node (false,
 [('o',
 Node (false,
 [('r',
 Node (false,
 [('r',
 Node (false,
 [('i',
 Node (false,
 [('f',
 Node (false,
 [('i', Node (false, [('c', Node (true, []))]))]))]));
 ('o', Node (false, [('r', Node (true, []))]))]))]))]))]))])

Answer 1

Sometimes just adding whitespace helps:

let rec add w tr = match w, tr with 
 | [], Node (_, lvl_l) ->
 Node (true, lvl_l)
 | wh :: wt, Node (b, lvl_l) -> 
 try
 let ins_point = List.assoc wh lvl_l in 
 Node (b, (wh, add wt ins_point)::(List.remove_assoc wh lvl_l)) 
 with Not_found ->
 Node (b, (wh, add wt empty)::lvl_l)

The match expression was upgraded in recent years to match on exceptions. So you can also do this:

let rec add w tr = match w, tr with 
 | [], Node (_, lvl_l) ->
 Node (true, lvl_l)
 | wh :: wt, Node (b, lvl_l) -> 
 begin match List.assoc wh lvl_l with
 | ins_point -> 
 Node (b, (wh, add wt ins_point)::(List.remove_assoc wh lvl_l)) 
 | exception Not_found ->
 Node (b, (wh, add wt empty)::lvl_l)
 end

(I had to wrap the match in a begin/end block to avoid confusing which match the patterns are applied.)

Since the return value has a common form, you could use match to just compute the second part of the pair:

let rec add w tr = match w, tr with 
 | [], Node (_, lvl_l) ->
 Node (true, lvl_l)
 | wh :: wt, Node (b, lvl_l) ->
 Node (b, match List.assoc wh lvl_l with
 | ins_point -> 
 (wh, add wt ins_point) :: (List.remove_assoc wh lvl_l) 
 | exception Not_found ->
 (wh, add wt empty) :: lvl_l)

It's subjective whether these are any better than your example. I do think, however, that visually showing the structure helps me understand the code better. For instance, my last version shows Node being produced in both match cases. You could elevate the constructor "one more level up" and have the nested match expressions compute the pair to apply to the constructor. This wasn't immediately obvious from your example.

Answer 2

Given that tr is only ever immediately pattern-matched on, and only has one constructor, you might immediately pattern-match it in the argument list.

E.g.

let rec add w (Node (b, lvl_l)) =
 match w with 
 | [] -> Node (true, lvl_l)
 | wh :: wt -> 
 try 
 let ins_point = List.assoc wh lvl_l in 
 Node (b, (wh, add wt ins_point) :: List.remove_assoc wh lvl_l) 
 with 
 Not_found -> Node (b, (wh, add wt empty) :: lvl_l)

The visual weight of your match is now much reduced.