Update markov chains to be backed by maps instead of lists

This should hopefully speed up lookups and stuff like that.

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>

lib/contrib/markov/chain.ex

@@ -2,7 +2,7 @@ defmodule Omnibot.Contrib.Markov.Chain do
   alias Omnibot.{Contrib.Markov.Chain, Util}
 
   @enforce_keys [:order]
-  defstruct order: 2, chain: []
+  defstruct order: 2, chain: %{}
 
   def train(chain, line) when is_binary(line) do
     train(chain, line |> String.split(~r/\s+/))
@@ -25,16 +25,11 @@ defmodule Omnibot.Contrib.Markov.Chain do
       raise(ArgumentError, message: "invalid key (length #{length(key)} vs. order #{order})")
     end
 
-    chain = case find_index(chain, key) do
-      # Insert weight
-      nil -> [{key, %{word => increment}} | chain]
-      # Update weight
-      index -> List.update_at(
-          chain,
-          index,
-          fn {key, mapping} -> {key, Map.update(mapping, word, increment, &(&1 + increment))} end
-      )
-    end
+    # %{
+    #   ["word1", "word2"] => %{"target" => weight}
+    # }
+    chain = Map.update(chain, key, %{word => increment},
+      fn weights -> Map.update(weights, word, increment, &(increment + &1)) end)
     %Chain{chain: chain, order: order}
   end
 

test/contrib/markov/chain_test.exs

@@ -5,66 +5,66 @@ defmodule MarkovChainTest do
   test "chain train works correctly" do
     chain = %Chain {order: 2}
             |> Chain.train(~w(foo bar baz))
-    assert chain.chain == [
-      {["bar", "baz"], %{nil => 1}},
-      {["foo", "bar"], %{"baz" => 1}},
-    ]
+    assert chain.chain == %{
+      ["bar", "baz"] => %{nil => 1},
+      ["foo", "bar"] => %{"baz" => 1},
+    }
 
     chain = chain |> Chain.train(~w(foo bar baz))
 
-    assert chain.chain == [
-      {["bar", "baz"], %{nil => 2}},
-      {["foo", "bar"], %{"baz" => 2}},
-    ]
+    assert chain.chain == %{
+      ["bar", "baz"] => %{nil => 2},
+      ["foo", "bar"] => %{"baz" => 2},
+    }
 
     chain = chain |> Chain.train(~w(baz bar foo))
 
-    assert chain.chain == [
-      {["bar", "foo"], %{nil => 1}},
-      {["baz", "bar"], %{"foo" => 1}},
-      {["bar", "baz"], %{nil => 2}},
-      {["foo", "bar"], %{"baz" => 2}},
-    ]
+    assert chain.chain == %{
+      ["bar", "foo"] => %{nil => 1},
+      ["baz", "bar"] => %{"foo" => 1},
+      ["bar", "baz"] => %{nil => 2},
+      ["foo", "bar"] => %{"baz" => 2},
+    }
 
     chain = chain |> Chain.train(~w(a b c))
-    assert chain.chain == [
-      {["b", "c"], %{nil => 1}},
-      {["a", "b"], %{"c" => 1}},
-      {["bar", "foo"], %{nil => 1}},
-      {["baz", "bar"], %{"foo" => 1}},
-      {["bar", "baz"], %{nil => 2}},
-      {["foo", "bar"], %{"baz" => 2}},
-    ]
+    assert chain.chain == %{
+      ["b", "c"] => %{nil => 1},
+      ["a", "b"] => %{"c" => 1},
+      ["bar", "foo"] => %{nil => 1},
+      ["baz", "bar"] => %{"foo" => 1},
+      ["bar", "baz"] => %{nil => 2},
+      ["foo", "bar"] => %{"baz" => 2},
+    }
   end
 
   test "chain add_weight works correctly" do
     chain = %Chain {order: 2}
             |> Chain.add_weight(["foo", "bar"], "baz")
-    assert chain.chain == [
-      {["foo", "bar"], %{"baz" => 1}}
-    ]
+    assert chain.chain == %{
+      ["foo", "bar"] => %{"baz" => 1},
+    }
 
     chain = chain |> Chain.add_weight(["foo", "bar"], "baz", 2)
 
-    assert chain.chain == [
-      {["foo", "bar"], %{"baz" => 3}}
-    ]
+    assert chain.chain == %{
+      ["foo", "bar"] => %{"baz" => 3},
+    }
 
     chain = chain |> Chain.add_weight(["foo", "bar"], "qux")
-    assert chain.chain == [
-      {["foo", "bar"], %{"baz" => 3, "qux" => 1}}
-    ]
+    assert chain.chain == %{
+      ["foo", "bar"] => %{"baz" => 3, "qux" => 1},
+    }
 
     chain = chain |> Chain.add_weight(["bar", "baz"], "qux")
-    assert chain.chain == [
-      {["bar", "baz"], %{"qux" => 1}},
-      {["foo", "bar"], %{"baz" => 3, "qux" => 1}},
-    ]
+    assert chain.chain == %{
+      ["bar", "baz"] => %{"qux" => 1},
+      ["foo", "bar"] => %{"baz" => 3, "qux" => 1},
+    }
 
     chain = chain |> Chain.add_weight(["bar", "baz"], nil)
-    assert chain.chain == [
-      {["bar", "baz"], %{"qux" => 1, nil => 1}},
-      {["foo", "bar"], %{"baz" => 3, "qux" => 1}},
-    ]
+    assert chain.chain == %{
+      ["bar", "baz"] => %{"qux" => 1, nil => 1},
+      ["foo", "bar"] => %{"baz" => 3, "qux" => 1},
+    }
   end
 end