Replace exception-driven Variant encoding

lib/ch/row_binary.ex

@@ -331,7 +331,7 @@ defmodule Ch.RowBinary do
   def encode({:variant, _types}, nil), do: 255
 
   def encode({:variant, types}, value) do
-    try_encode_variant(types, 0, value)
+    encode_variant(types, 0, value)
   end
 
   def encode(:datetime, %NaiveDateTime{} = datetime) do
@@ -493,21 +493,116 @@ defmodule Ch.RowBinary do
 
   defp encode_many_kv([] = done, _key_type, _value_type), do: done
 
-  # TODO find a better way than try/rescue
-  defp try_encode_variant([type | types], idx, value) do
-    try do
-      encode(type, value)
+  defp encode_variant([type | types], idx, value) do
+    if variant_matches?(type, value) do
+      [idx | encode(type, value)]
     else
-      encoded -> [idx | encoded]
-    rescue
-      _e -> try_encode_variant(types, idx + 1, value)
+      encode_variant(types, idx + 1, value)
     end
   end
 
-  defp try_encode_variant([], _idx, value) do
+  defp encode_variant([], _idx, value) do
     raise ArgumentError, "no matching type found for encoding #{inspect(value)} as Variant"
   end
 
+  defp variant_matches?(:string, value), do: is_binary(value) or is_list(value)
+  defp variant_matches?({:fixed_string, _size}, value), do: is_binary(value)
+  defp variant_matches?(:boolean, value), do: is_boolean(value)
+
+  defp variant_matches?(type, value)
+       when type in [:u8, :u16, :u32, :u64, :u128, :u256],
+       do: is_integer(value) and value >= 0
+
+  defp variant_matches?(type, value)
+       when type in [:i8, :i16, :i32, :i64, :i128, :i256],
+       do: is_integer(value)
+
+  defp variant_matches?(type, value) when type in [:f32, :f64], do: is_number(value)
+
+  defp variant_matches?({type, _scale}, value)
+       when type in [:decimal32, :decimal64, :decimal128, :decimal256],
+       do: is_struct(value, Decimal)
+
+  defp variant_matches?({:array, type}, value) when is_list(value) do
+    variant_list_matches?(value, type)
+  end
+
+  defp variant_matches?({:tuple, types}, value) when is_tuple(value) do
+    variant_tuple_matches?(types, Tuple.to_list(value))
+  end
+
+  defp variant_matches?({:tuple, types}, value) when is_list(value) do
+    variant_tuple_matches?(types, value)
+  end
+
+  defp variant_matches?({:map, key_type, value_type}, value) when is_map(value) do
+    variant_map_matches?(Map.to_list(value), key_type, value_type)
+  end
+
+  defp variant_matches?({:map, key_type, value_type}, value) when is_list(value) do
+    variant_map_matches?(value, key_type, value_type)
+  end
+
+  defp variant_matches?({:nullable, _type}, nil), do: true
+  defp variant_matches?({:nullable, type}, value), do: variant_matches?(type, value)
+  defp variant_matches?(:date, value), do: is_struct(value, Date)
+  defp variant_matches?(:date32, value), do: is_struct(value, Date)
+
+  defp variant_matches?(:datetime, value) do
+    is_struct(value, NaiveDateTime) or is_struct(value, DateTime)
+  end
+
+  defp variant_matches?({:datetime64, _time_unit}, value) do
+    is_struct(value, NaiveDateTime) or is_struct(value, DateTime)
+  end
+
+  defp variant_matches?(:time, value), do: is_struct(value, Time)
+  defp variant_matches?({:time64, _time_unit}, value), do: is_struct(value, Time)
+  defp variant_matches?(:uuid, value), do: is_binary(value)
+  defp variant_matches?(:ipv4, value), do: is_tuple(value) and tuple_size(value) == 4
+
+  defp variant_matches?(:ipv6, value) do
+    (is_tuple(value) and tuple_size(value) == 8) or
+      (is_binary(value) and byte_size(value) == 16)
+  end
+
+  defp variant_matches?(:point, value) do
+    is_tuple(value) and tuple_size(value) == 2
+  end
+
+  defp variant_matches?({type, _mapping}, value) when type in [:enum8, :enum16] do
+    is_integer(value) or is_binary(value)
+  end
+
+  defp variant_matches?(:dynamic, value) do
+    is_binary(value) or is_integer(value) or is_float(value) or is_struct(value, Date) or
+      is_struct(value, NaiveDateTime) or value == []
+  end
+
+  defp variant_matches?(:json, _value), do: true
+  defp variant_matches?(_type, _value), do: false
+
+  defp variant_list_matches?([value | values], type) do
+    variant_matches?(type, value) and variant_list_matches?(values, type)
+  end
+
+  defp variant_list_matches?([], _type), do: true
+
+  defp variant_tuple_matches?([type | types], [value | values]) do
+    variant_matches?(type, value) and variant_tuple_matches?(types, values)
+  end
+
+  defp variant_tuple_matches?([], []), do: true
+  defp variant_tuple_matches?(_types, _values), do: false
+
+  defp variant_map_matches?([{key, value} | values], key_type, value_type) do
+    variant_matches?(key_type, key) and variant_matches?(value_type, value) and
+      variant_map_matches?(values, key_type, value_type)
+  end
+
+  defp variant_map_matches?([], _key_type, _value_type), do: true
+  defp variant_map_matches?(_value, _key_type, _value_type), do: false
+
   @compile {:inline, d: 1}
 
   defp d(?0), do: 0

test/ch/row_binary_test.exs

@@ -233,6 +233,29 @@ defmodule Ch.RowBinaryTest do
         encode({:variant, [:u8]}, "x")
       end
     end
+
+    test "variant selects the first recursively compatible type" do
+      assert encode({:variant, [:string, :u64]}, 42) == [1 | <<42::64-little>>]
+
+      assert encode({:variant, [{:array, :u64}, {:array, :string}]}, ["x"]) ==
+               [1, 1, [1 | "x"]]
+
+      assert encode({:variant, [{:tuple, [:string, :u64]}, {:tuple, [:u64, :string]}]}, {1, "x"}) ==
+               [1, <<1::64-little>>, [1 | "x"]]
+
+      assert encode(
+               {:variant, [{:map, :string, :u64}, {:map, :string, {:nullable, :string}}]},
+               %{"x" => nil}
+             ) == [1, 1, [1 | "x"], 1]
+    end
+
+    test "variant does not hide errors from a compatible encoder" do
+      type = {:variant, [{:enum8, %{"present" => 1}}, :string]}
+
+      assert_raise ArgumentError,
+                   ~s[enum value "missing" not found in mapping: %{"present" => 1}],
+                   fn -> encode(type, "missing") end
+    end
   end
 
   test "strings preserve raw bytes" do