Rewrite VariableReference, and support iterating through any type

lib/orb/dsl/dsl.ex

@@ -357,10 +357,21 @@ defmodule Orb.DSL do
                   nil
 
                 identifier ->
+                  source_type =
+                    quote do
+                      case unquote(source) do
+                        %Range{} ->
+                          Orb.I32
+
+                        %Orb.VariableReference{push_type: source_iterator} ->
+                          source_iterator.value_type()
+                      end
+                    end
+
                   quote(
                     do:
                       var!(unquote(var)) =
-                        Orb.VariableReference.local(unquote(identifier), Orb.I32)
+                        Orb.VariableReference.local(unquote(identifier), unquote(source_type))
                   )
               end,
             identifier: identifier,
@@ -370,7 +381,6 @@ defmodule Orb.DSL do
         %Range{first: first, last: last, step: 1} ->
           with do
             element_type = Orb.I32
-
             _ = init_elixir_var
 
             Orb.InstructionSequence.new([
@@ -405,9 +415,9 @@ defmodule Orb.DSL do
             |> Orb.InstructionSequence.concat_locals([{identifier, element_type}])
           end
 
-        source = %{push_type: source_iterator} when not is_nil(source_iterator) ->
+        %Orb.VariableReference{push_type: source_iterator} = source ->
           with do
-            element_type = Orb.I32
+            element_type = source_iterator.value_type()
             _ = init_elixir_var
 
             body =

lib/orb/i64.ex

@@ -12,6 +12,7 @@ defmodule Orb.I64 do
   end
 
   for op <- Ops.i64(1) do
+    @doc Ops.doc(:i64, op)
     def unquote(op)(a) do
       Instruction.i64(unquote(op), a)
     end

lib/orb/instruction_sequence.ex

@@ -66,6 +66,7 @@ defmodule Orb.InstructionSequence do
   def new(push_type, instructions, opts \\ []) when is_list(instructions) do
     %__MODULE__{
       push_type: push_type,
+      pop_type: Keyword.get(opts, :pop_type, nil),
       body: instructions,
       locals: Keyword.get(opts, :locals, [])
     }

lib/orb/iterator.ex

@@ -11,7 +11,12 @@ defmodule Orb.Iterator do
   @callback valid?(var_ref) :: %Orb.Instruction{} | integer()
 
   @doc """
-  Extracts or calculates the current value.
+  The type for iterated elements.
+  """
+  @callback value_type() :: atom() | tuple()
+
+  @doc """
+  Extracts or calculates the current element value.
   """
   @callback value(var_ref) :: %Orb.Instruction{} | integer() | float()
 

lib/orb/ops.ex

@@ -256,4 +256,9 @@ defmodule Orb.Ops do
   defp nth(1), do: "2nd"
   defp nth(2), do: "3rd"
   defp nth(n), do: "#{n + 1}th"
+
+  def doc(type, op)
+  def doc(:i64, :extend_i32_u), do: "Convert unsigned i32 to unsigned i64."
+  def doc(:i64, :extend_i32_s), do: "Convert signed i32 to signed i64."
+  def doc(_type, _op), do: nil
 end

lib/orb/variable_reference.ex

@@ -1,76 +1,136 @@
 defmodule Orb.VariableReference do
   @moduledoc false
 
-  defstruct [:global_or_local, :identifier, :push_type]
+  defmodule Global do
+    defstruct [:identifier, :type]
 
-  alias Orb.Instruction
+    def set(%__MODULE__{identifier: identifier, type: type}, new_value) do
+      Orb.Instruction.Global.Set.new(type, identifier, new_value)
+    end
+
+    def set_from_stack(%__MODULE__{identifier: identifier, type: type}) do
+      Orb.Instruction.Global.Set.new(type, identifier, Orb.Stack.Pop.new(type))
+    end
+
+    defimpl Orb.ToWat do
+      def to_wat(%Orb.VariableReference.Global{identifier: identifier}, indent) do
+        [indent, "(global.get $", to_string(identifier), ?)]
+      end
+    end
+  end
+
+  defmodule Local do
+    defstruct [:identifier, :type]
+
+    def set(%__MODULE__{identifier: identifier, type: type}, new_value) do
+      Orb.Instruction.local_set(type, identifier, new_value)
+    end
+
+    def set_from_stack(%__MODULE__{identifier: identifier, type: type}) do
+      Orb.Instruction.local_set(type, identifier, Orb.Stack.Pop.new(type))
+    end
+
+    defimpl Orb.ToWat do
+      def to_wat(%Orb.VariableReference.Local{identifier: identifier}, indent) do
+        [indent, "(local.get $", to_string(identifier), ?)]
+      end
+    end
+
+    defimpl Orb.ToWasm do
+      import Orb.Leb
+
+      def to_wasm(%Orb.VariableReference.Local{identifier: identifier}, context) do
+        import Orb.Leb
+
+        [0x20, leb128_u(Orb.ToWasm.Context.fetch_local_index!(context, identifier))]
+      end
+    end
+  end
+
+  defstruct [:entries, :push_type, :pop_type]
 
   def global(identifier, type) do
-    %__MODULE__{global_or_local: :global, identifier: identifier, push_type: type}
+    %__MODULE__{
+      entries: [%Global{identifier: identifier, type: type}],
+      push_type: type
+    }
   end
 
   def local(identifier, type) do
-    %__MODULE__{global_or_local: :local, identifier: identifier, push_type: type}
+    %__MODULE__{
+      entries: [%Local{identifier: identifier, type: type}],
+      push_type: type
+    }
   end
 
   def set(
-        %__MODULE__{global_or_local: :local, identifier: identifier, push_type: type},
+        %__MODULE__{entries: entries},
         new_value
       ) do
-    Instruction.local_set(type, identifier, new_value)
-  end
+    new_value =
+      case new_value do
+        new_value when is_tuple(new_value) ->
+          Tuple.to_list(new_value)
 
-  def as_set(%__MODULE__{global_or_local: :local, identifier: identifier, push_type: type}) do
-    Instruction.local_set(type, identifier)
-  end
+        new_value when is_list(new_value) ->
+          new_value
 
-  @behaviour Access
+        new_value ->
+          List.wrap(new_value)
+      end
 
-  @impl Access
-  # TODO: I think this should only live on custom types like UnsafePointer
-  def fetch(
-        %__MODULE__{global_or_local: :local, identifier: _identifier, push_type: :i32} = ref,
-        at: offset
-      ) do
-    ast = Instruction.i32(:load, Instruction.i32(:add, ref, offset))
-    {:ok, ast}
+    for {entry, value} <- Enum.zip(entries, new_value) do
+      entry.__struct__.set(entry, value)
+    end
+    |> Orb.InstructionSequence.new()
   end
 
-  def fetch(
-        %__MODULE__{global_or_local: :local, identifier: _identifier, push_type: mod} = ref,
-        key
-      ) do
-    mod.fetch(ref, key)
+  def as_set(%__MODULE__{entries: entries}) when length(entries) === 1 do
+    Orb.InstructionSequence.new(
+      nil,
+      for entry <- entries do
+        entry.__struct__.set_from_stack(entry)
+      end,
+      pop_type: hd(entries).type
+    )
   end
 
-  @impl Access
-  def get_and_update(_data, _key, _function) do
-    raise UndefinedFunctionError, module: __MODULE__, function: :get_and_update, arity: 3
-  end
+  with @behaviour Access do
+    @impl Access
+    def fetch(
+          %__MODULE__{entries: [%Local{type: mod}]} = var_ref,
+          key
+        ) do
+      mod.fetch(var_ref, key)
+    end
 
-  @impl Access
-  def pop(_data, _key) do
-    raise UndefinedFunctionError, module: __MODULE__, function: :pop, arity: 2
-  end
+    @impl Access
+    def get_and_update(_data, _key, _function) do
+      raise UndefinedFunctionError, module: __MODULE__, function: :get_and_update, arity: 3
+    end
 
-  defimpl Orb.ToWat do
-    def to_wat(%Orb.VariableReference{global_or_local: :global, identifier: identifier}, indent) do
-      [indent, "(global.get $", to_string(identifier), ?)]
+    @impl Access
+    def pop(_data, _key) do
+      raise UndefinedFunctionError, module: __MODULE__, function: :pop, arity: 2
     end
+  end
 
-    def to_wat(%Orb.VariableReference{global_or_local: :local, identifier: identifier}, indent) do
-      [indent, "(local.get $", to_string(identifier), ?)]
+  defimpl Orb.ToWat do
+    def to_wat(%Orb.VariableReference{entries: entries}, indent) do
+      for entry <- entries do
+        Orb.ToWat.to_wat(entry, indent)
+      end
     end
   end
 
   defimpl Orb.ToWasm do
-    import Orb.Leb
-
     def to_wasm(
-          %Orb.VariableReference{global_or_local: :local, identifier: identifier},
+          %Orb.VariableReference{entries: entries},
           context
         ) do
-      [0x20, leb128_u(Orb.ToWasm.Context.fetch_local_index!(context, identifier))]
+      for entry <- entries do
+        Orb.ToWasm.to_wasm(entry, context)
+      end
     end
   end
 end

test/loop_test.exs

@@ -157,12 +157,34 @@ defmodule LoopTest do
         @impl b
         def valid?(var), do: Orb.I32.le_u(var, ?z)
         @impl b
+        def value_type(), do: Orb.I32
+        @impl b
         def value(var), do: var
         @impl b
         def next(var), do: Orb.I32.add(var, 1)
       end
     end
 
+    defmodule AlphabetIterator64 do
+      def new(), do: ?a
+
+      with @behaviour b = Orb.CustomType do
+        @impl b
+        defdelegate wasm_type, to: Orb.I32
+      end
+
+      with @behaviour b = Orb.Iterator do
+        @impl b
+        def valid?(var), do: Orb.I32.le_u(var, ?z)
+        @impl b
+        def value_type(), do: Orb.I64
+        @impl b
+        def value(var), do: Orb.I64.extend_i32_u(var)
+        @impl b
+        def next(var), do: Orb.I32.add(var, 1)
+      end
+    end
+
     defmodule IteratorConsumer do
       use Orb
 
@@ -195,12 +217,23 @@ defmodule LoopTest do
 
         i
       end
+
+      defw test64, I64, i: I64, alpha: AlphabetIterator64 do
+        alpha = AlphabetIterator64.new()
+
+        loop char <- alpha do
+          i = i + char
+        end
+
+        i
+      end
     end
 
     test "behaves correctly" do
       assert 26 = Wasm.call(IteratorConsumer, :test)
       assert 26 = Wasm.call(IteratorConsumer, :test2)
       assert 2847 = Wasm.call(IteratorConsumer, :test3)
+      assert 2847 = Wasm.call(IteratorConsumer, :test64)
       assert 2847 = Enum.sum(?a..?z)
     end
 
@@ -267,6 +300,27 @@ defmodule LoopTest do
                    )    )
                  (local.get $i)
                )
+               (func $test64 (export "test64") (result i64)
+                 (local $i i64)
+                 (local $alpha i32)
+                 (local $char i64)
+                 (i32.const 97)
+                 (local.set $alpha)
+                 (loop $char
+                   (i32.le_u (local.get $alpha) (i32.const 122))
+                   (if
+                     (then
+                       (i64.extend_i32_u (local.get $alpha))
+                       (local.set $char)
+                       (i64.add (local.get $i) (local.get $char))
+                       (local.set $i)
+                       (i32.add (local.get $alpha) (i32.const 1))
+                       (local.set $alpha)
+                       (br $char)
+                     )
+                   )    )
+                 (local.get $i)
+               )
              )
              """ = Orb.to_wat(IteratorConsumer)
     end