feat(Trie): completion API

src/Trie.ml

@@ -344,3 +344,39 @@ let set_of_tags (type tag) (cmp : tag -> tag -> int) (v : ('data, tag) t) : tag
   let set = ref TagSet.empty in
   Option.iter (fun (_, n) -> iter_tag_node (fun t -> set := TagSet.add t !set) n) v;
   TagSet.to_seq !set
+
+let edit_distance ~cutoff x y =
+  let len_x, len_y = String.length x, String.length y in
+  let grid = Array.make_matrix (len_x + 1) (len_y + 1) 0 in
+  for i = 1 to len_x do
+    grid.(i).(0) <- i;
+  done;
+  for j = 1 to len_y do
+    grid.(0).(j) <- j;
+  done;
+  for j = 1 to len_y do
+    for i = 1 to len_x do
+      let cost = if x.[i-1] = y.[j-1] then 0 else 1 in
+      let k = Int.min (grid.(i-1).(j) + 1) (grid.(i).(j-1) + 1)in
+      grid.(i).(j) <- Int.min k (grid.(i-1).(j-1) + cost)
+    done;
+  done;
+  let result = grid.(len_x).(len_y) in
+  if result > cutoff
+    then None
+  else
+    Some result
+
+let complete ?prefix ~(cutoff : int) (p : bwd_path) : ('data, 'tag) t -> ('data, int) t =
+  let compare p d = 
+    edit_distance ~cutoff (String.concat "" (Bwd.to_list p)) (String.concat "" (Bwd.to_list d))
+  in
+  filter_map ?prefix (fun q (data, _) -> 
+    match compare p q with
+    | Some i -> 
+      if i > cutoff then 
+        None 
+      else 
+        (Some (data, i))
+    | None -> None
+  )

src/Trie.mli

@@ -78,6 +78,12 @@ val filter : ?prefix:bwd_path -> (bwd_path -> 'data * 'tag -> bool) -> ('data, '
 *)
 val filter_map : ?prefix:bwd_path -> (bwd_path -> 'data1 * 'tag1 -> ('data2 * 'tag2) option) -> ('data1, 'tag1) t -> ('data2, 'tag2) t
 
+(** [complete ~cutoff:i p trie] retags each entry [e] in [trie] with the edit distance of p to the path of [e]. It can be used to implement autocomplete and "Did you mean..." style diagnostics.*)
+val complete : ?prefix:bwd_path -> cutoff:int -> bwd_path -> ('data, 'tag) t -> ('data, int) t
+
+(** A simple implementation of the Levenshtein edit distance algorithm, used in {!val:complete} complete*)
+val edit_distance : cutoff:int -> string -> string -> int option
+
 (** {1 Updating} *)
 
 (** [update_subtree p f t] replaces the subtree [t'] rooted at [p] in [t] with [f t']. *)

test/Example.expected

@@ -1,27 +1,40 @@
 [Info] Got the following bindings at (root):
+  foo => 1 (local)
   x => 1 (local)
 
 [Warning] Data 1 (local) assigned at x was shadowed by data 2 (local) in the export namespace.
 [Warning] Data 1 (local) assigned at x was shadowed by data 2 (local) in the visible namespace.
 [Info] Got the following bindings at (root):
+  foo => 1 (local)
   x => 2 (local)
 
 [Info] Got the following bindings at (root):
+  foo => 1 (local)
   x => 10 (local)
 
 [Info] Got the following bindings at (root):
+  foo => 1 (local)
   x => 10 (local)
   z.y => 20 (imported)
 
 [Info] Got the following bindings at (root):
   a => 100 (local)
+  foo => 1 (local)
   x => 10 (local)
   z.y => 20 (imported)
 
 [Warning] Data 10 (local) assigned at w.x was shadowed by data 10 (local).
 [Info] Got the following bindings at (root):
+  foo => 1 (local)
   w.a => 100 (local)
   w.x => 10 (local)
   x => 10 (local)
   z.y => 20 (imported)
 
+[Info] Got the following completion items for path zoo:
+  foo, distance: 1
+  fooo, distance: 2
+  goo, distance: 1
+  ogoo, distance: 2
+  z.y, distance: 2
+

test/Example.ml

@@ -2,7 +2,7 @@ open Yuujinchou
 open Bwd
 
 (* A tiny language demonstrating some power of the Scope module. *)
-type modifier_cmd = Print
+type modifier_cmd = Print | Complete
 type decl =
   (* declaration *)
   | Decl of Trie.path * int
@@ -12,6 +12,8 @@ type decl =
   | Import of int Trie.Untagged.t * modifier_cmd Language.t
   (* printing out all visible bindings *)
   | PrintVisible
+  (* Get completion suggestions for a path *)
+  | CompleteVisible
   (* exporting a binding *)
   | Export of Trie.path
   (* section *)
@@ -74,6 +76,19 @@ struct
         input;
       Format.printf "@]@.";
       input
+    | Complete ->
+      let typo = Bwd.of_list ["zoo"] in
+      Format.printf "@[<v 2>[Info] Got the following completion items for path %a%a:@;"
+        pp_path typo pp_context context;
+      List.iter
+        (fun (path, dist) -> Format.printf "%a, distance: %i@;" pp_path path dist)
+        (Trie.complete ~cutoff:2 typo input
+          |> Trie.to_seq
+          |> List.of_seq
+          |> List.map (fun (path, (_, dist)) -> (Bwd.of_list path), dist)
+        );
+      Format.printf "@]@.";
+      input
 end
 
 (* The interpreter *)
@@ -89,6 +104,8 @@ let rec interpret_decl : decl -> unit =
     S.import_subtree ~modifier:m ([], t)
   | PrintVisible ->
     S.modify_visible (Language.hook Print)
+  | CompleteVisible ->
+    S.modify_visible (Language.hook Complete)
   | Export p ->
     S.export_visible (Language.only p)
   | Section (p, sec) ->
@@ -105,6 +122,7 @@ let interpret (prog : program) =
 (* Some code in action *)
 let () = interpret [
     Decl (["x"], 1);
+    Decl (["foo"], 1);
     PrintVisible;
     Decl (["x"], 2);
     PrintVisible;
@@ -120,4 +138,8 @@ let () = interpret [
         Export ["x"];
       ]);
     PrintVisible;
+    Decl (["fooo"], 1);
+    Decl (["goo"], 1);
+    Decl (["ogoo"], 1);
+    CompleteVisible;
   ]

test/ListAsTrie.ml

@@ -106,3 +106,19 @@ let retag_subtree pre t l =
   List.map (fun ((p, (d, _)) as b) -> if Option.is_some (split_path pre p) then p, (d, t) else b) l
 let untag l = retag () l
 let set_of_tags cmp l = List.to_seq @@ List.sort_uniq cmp @@ List.map (fun (_, (_, t)) -> t) l
+
+let edit_distance = Yuujinchou.Trie.edit_distance
+
+let complete ?prefix ~(cutoff : int) (p : bwd_path) : ('data, 'tag) t -> ('data, int) t =
+  let compare p d = 
+    edit_distance ~cutoff (String.concat "" (Bwd.to_list p)) (String.concat "" (Bwd.to_list d))
+  in
+  filter_map ?prefix (fun q (data, _) -> 
+    match compare p q with
+    | Some i -> 
+      if i > cutoff then 
+        None 
+      else 
+        (Some (data, i))
+    | None -> None
+  )

test/TestTrie.ml

@@ -248,6 +248,10 @@ let test_untag =
 let test_set_of_tags =
   Q.Test.make ~count ~name:"set_of_tags" gen_list ~print:print_list
     (fun l -> List.of_seq (Trie.set_of_tags Int.compare (of_list l)) = List.of_seq (ListAsTrie.set_of_tags Int.compare l))
+let test_complete =
+  Q.Test.make ~count ~name:"complete" Q.Gen.(pair gen_bwd_path gen_list) ~print:Q.Print.(pair print_bwd_path print_list)
+    (fun (p, l) -> 
+      (to_list @@ Trie.complete ~cutoff:2 p (of_list l)) = ListAsTrie.complete ~cutoff:2 p l)
 
 let () =
   exit @@
@@ -285,4 +289,5 @@ let () =
     ; test_retag_subtree
     ; test_untag
     ; test_set_of_tags
+    ; test_complete
     ]