1 files changed, 43 insertions, 32 deletions
diff --git a/motus/lib/entropy.ml b/motus/lib/entropy.ml
index 043d0c8..aec5d86 100644
--- a/motus/lib/entropy.ml
+++ b/motus/lib/entropy.ml
@@ -43,41 +43,52 @@ let get_entropy max_element words_number arr =
 
 let analyse : int -> catalog:Wordlist.t -> Wordlist.t -> t =
  fun base ~catalog words ->
-  let max_element = Float.to_int @@ (Validity.elements ** Float.of_int base) in
-  let words_number = Float.of_int (Wordlist.list_size catalog) in
+  (* If we have only two elements, just pick one of them *)
+  if Wordlist.list_size words <= 2
+  then (1.0, Option.get @@ Wordlist.pick words)
+  else
+    let max_element =
+      Float.to_int @@ (Validity.elements ** Float.of_int base)
+    in
+    let words_number = Float.of_int (Wordlist.list_size catalog) in
 
-  let arr = Bigarray.Array1.create Bigarray.Int Bigarray.C_layout max_element in
-  match Wordlist.pick words with
-  | None -> (0., "")
-  | Some v ->
-      (* Build the array  *)
-      Seq.fold_left
-        (fun (score, word) word_ref ->
-          (* Reinitialize the array (we use the same in the successive
-             iterations *)
-          let set_ref = String.to_seq word_ref |> CharSet.of_seq in
+    let arr =
+      Bigarray.Array1.create Bigarray.Int Bigarray.C_layout max_element
+    in
+    match Wordlist.pick words with
+    | None -> (0., "")
+    | Some v ->
+        (* Build the array  *)
+        Seq.fold_left
+          (fun (score, word) word_ref ->
+            (* Reinitialize the array (we use the same in the successive
+               iterations *)
+            let set_ref = String.to_seq word_ref |> CharSet.of_seq in
 
-          Bigarray.Array1.fill arr 0;
+            Bigarray.Array1.fill arr 0;
 
-          Seq.iter
-            (fun w2 ->
-              let result = compare_words ~ref:(word_ref, set_ref) w2 in
-              match result with
-              | None -> ()
-              | Some r ->
-                  let idx = Validity.index_of_result r in
+            Seq.iter
+              (fun w2 ->
+                let result = compare_words ~ref:(word_ref, set_ref) w2 in
+                match result with
+                | None -> ()
+                | Some r ->
+                    let idx = Validity.index_of_result r in
 
-                  let content = Bigarray.Array1.get arr idx in
-                  Bigarray.Array1.set arr idx (succ content) )
-            (Wordlist.words words);
+                    let content = Bigarray.Array1.get arr idx in
+                    Bigarray.Array1.set arr idx (succ content) )
+              (Wordlist.words words);
 
-          (* Now evaluate the entropy in the array *)
-          let entropy = get_entropy max_element words_number arr in
+            (* Now evaluate the entropy in the array *)
+            let entropy = get_entropy max_element words_number arr in
 
-          if !entropy > score
-          then (
-            Printf.printf "Entropy for selecting %s : %.2f\n" word_ref !entropy;
-            (!entropy, word_ref) )
-          else (score, word) )
-        (-0., v)
-        (Wordlist.words catalog)
+            if !entropy > score
+            then (
+              Printf.printf
+                "Entropy for selecting %s : %.2f\n"
+                word_ref
+                !entropy;
+              (!entropy, word_ref) )
+            else (score, word) )
+          (-0., v)
+          (Wordlist.words catalog)