1 files changed, 24 insertions, 17 deletions

diff --git a/motus/lib/validity.ml b/motus/lib/validity.ml
index ae85f04..0fdc40c 100644
--- a/motus/lib/validity.ml
+++ b/motus/lib/validity.ml
@@ -12,25 +12,28 @@ type t =
 
 [@@@warning "+32"]
 
-let m = Float.of_int (1 + max)
+let elements = Float.of_int (1 + max)
 
 (** Get the index of a validity result *)
 let index_of_result : t array -> int =
  fun elems ->
-  let _, value =
-    Array.fold_left elems ~init:(0., 0.) ~f:(fun (pos, acc) content ->
-        let v = Float.of_int (to_enum content) in
-        let acc' = acc +. (v *. (m ** pos)) in
-        (pos +. Float.one, acc') )
-  in
-  Float.to_int value
+  (* This seems to be more efficient than a foldLeft *)
+  let value = ref 0. in
+  Array.iteri
+    ~f:(fun pos content ->
+      let pos = Float.of_int pos in
+      let v = Float.of_int (to_enum content) in
+      let acc' = !value +. (v *. (elements ** pos)) in
+      value := acc' )
+    elems;
+  Float.to_int !value
 
 
 let index_to_result : base:int -> int -> t array =
  fun ~base n ->
   let rec _f acc n i =
-    let next = Float.round (n /. m)
-    and rem = Float.(to_int @@ rem n m) in
+    let next = Float.round (n /. elements)
+    and rem = Float.(to_int @@ rem n elements) in
 
     match (rem, i) with
     | _, 0 -> Array.of_list acc
@@ -45,29 +48,33 @@ let index_to_result : base:int -> int -> t array =
 (** Build a sequence of all the possible status for a given number of letters *)
 let sequence : int -> t array Seq.t =
  fun base ->
-  let max_element = Float.to_int @@ (m ** Float.of_int base) in
+  let max_element = Float.to_int @@ (elements ** Float.of_int base) in
 
   Seq.unfold
     (fun n ->
-      if n < max_element then Some (index_to_result ~base n, n + 1) else None )
+      if n < max_element then Some (index_to_result ~base n, succ n) else None
+      )
     0
 
 
-let compare_words : string -> ref:string -> t array option =
+module CharSet = Set.Make (Char)
+
+let compare_words : string -> ref:string * CharSet.t -> t array option =
  fun w1 ~ref ->
+  let wordRef = fst ref in
   let l1 = String.length w1 in
-  if l1 <> String.length ref
+  if l1 <> String.length wordRef
   then None
   else
     let result =
       Array.init l1 ~f:(fun i ->
           let c1 = String.get w1 i
-          and c2 = String.get ref i in
+          and c2 = String.get wordRef i in
 
           let state =
             if Char.equal c1 c2
             then Wellplaced
-            else if String.contains ref c1
+            else if CharSet.mem c2 (snd ref)
             then Misplaced
             else Missing
           in
@@ -92,7 +99,7 @@ let to_criterias : string -> t array -> Criteria.t list =
   let l, _ =
     Array.fold_left t ~init:([], 0) ~f:(fun (acc, i) t ->
         let acc = to_criteria (String.get word i) i t acc in
-        (acc, i + 1) )
+        (acc, succ i) )
   in
 
   List.rev l