Move expansion of Charms.Defm.if to Kernel.if

bench/enif_merge_sort.ex

@@ -40,7 +40,7 @@ defmodule ENIFMergeSort do
       left_term = Pointer.load(Term.t(), Pointer.element_ptr(Term.t(), left_temp, i))
       right_term = Pointer.load(Term.t(), Pointer.element_ptr(Term.t(), right_temp, j))
 
-      struct_if(enif_compare(left_term, right_term) <= 0) do
+      if(enif_compare(left_term, right_term) <= 0) do
         Pointer.store(
           Pointer.load(Term.t(), Pointer.element_ptr(Term.t(), left_temp, i)),
           Pointer.element_ptr(Term.t(), arr, k)
@@ -89,7 +89,7 @@ defmodule ENIFMergeSort do
   end
 
   defm do_sort(arr :: Pointer.t(), l :: i32(), r :: i32()) do
-    struct_if(l < r) do
+    if(l < r) do
       two_const = op arith.constant(value: Attribute.integer(i32(), 2)) :: i32()
       two = result_at(two_const, 0)
       m = op arith.divsi(l + r, two) :: i32()

bench/enif_quick_sort.ex

@@ -22,7 +22,7 @@ defmodule ENIFQuickSort do
     start = Pointer.element_ptr(Term.t(), arr, low)
 
     for_loop {element, j} <- {Term.t(), start, high - low} do
-      struct_if(enif_compare(element, pivot) < 0) do
+      if(enif_compare(element, pivot) < 0) do
         i = Pointer.load(i32(), i_ptr) + 1
         Pointer.store(i, i_ptr)
         j = value index.casts(j) :: i32()
@@ -40,7 +40,7 @@ defmodule ENIFQuickSort do
   end
 
   defm do_sort(arr :: Pointer.t(), low :: i32(), high :: i32()) do
-    struct_if(low < high) do
+    if(low < high) do
       pi = call partition(arr, low, high) :: i32()
       do_sort(arr, low, pi - 1)
       do_sort(arr, pi + 1, high)

bench/enif_tim_sort.ex

@@ -66,8 +66,8 @@ defmodule ENIFTimSort do
         right = op arith.minsi(left + 2 * size - 1, n - 1) :: i32()
         right = result_at(right, 0)
 
-        struct_if(mid < right) do
-          call ENIFMergeSort.merge(arr, left, mid, right) :: []
+        if(mid < right) do
+          call ENIFMergeSort.merge(arr, left, mid, right)
         end
 
         Pointer.store(left + 2 * size, left_ptr)

lib/charms/defm.ex

@@ -46,11 +46,6 @@ defmodule Charms.Defm do
   """
   defmacro cond_br(_condition, _clauses), do: :implemented_in_expander
 
-  @doc """
-  `if` expression requires identical types for both branches
-  """
-  defmacro struct_if(_condition, _clauses), do: :implemented_in_expander
-
   @doc """
   define a function that can be JIT compiled
 

lib/charms/defm/expander.ex

@@ -711,7 +711,7 @@ defmodule Charms.Defm.Expander do
     {v, state, env}
   end
 
-  defp expand_macro(_meta, Charms.Defm, :struct_if, [condition, clauses], _callback, state, env) do
+  defp expand_macro(_meta, Kernel, :if, [condition, clauses], _callback, state, env) do
     true_body = Keyword.fetch!(clauses, :do)
     false_body = clauses[:else]
     {condition, state, env} = expand(condition, state, env)
@@ -720,24 +720,50 @@ defmodule Charms.Defm.Expander do
       mlir ctx: state.mlir.ctx, block: state.mlir.blk do
         alias Beaver.MLIR.Dialect.SCF
 
+        b =
+          block _true() do
+            ret_t =
+              with {ret, _, _} <-
+                     expand(true_body, put_in(state.mlir.blk, Beaver.Env.block()), env),
+                   %MLIR.Value{} = ret when not is_nil(false_body) <-
+                     ret |> List.wrap() |> List.last() do
+                SCF.yield(ret) >>> []
+                MLIR.Value.type(ret)
+              else
+                %MLIR.Operation{} ->
+                  SCF.yield() >>> []
+                  []
+
+                %MLIR.Value{} = ret ->
+                  SCF.yield(ret) >>> []
+                  MLIR.Value.type(ret)
+              end
+          end
+
+        # TODO: doc about an expression which is a value and an operation
         SCF.if [condition] do
           region do
-            block _true() do
-              expand(true_body, put_in(state.mlir.blk, Beaver.Env.block()), env)
-              SCF.yield() >>> []
-            end
+            MLIR.CAPI.mlirRegionAppendOwnedBlock(Beaver.Env.region(), b)
           end
 
           region do
             block _false() do
-              if false_body do
-                expand(false_body, put_in(state.mlir.blk, Beaver.Env.block()), env)
+              with {ret, _, _} <-
+                     unless(is_nil(false_body),
+                       do: expand(false_body, put_in(state.mlir.blk, Beaver.Env.block()), env)
+                     ),
+                   %MLIR.Value{} = ret <- ret |> List.wrap() |> List.last() do
+                SCF.yield(ret) >>> []
+              else
+                %MLIR.Value{} = ret ->
+                  SCF.yield(ret) >>> []
+
+                _ ->
+                  SCF.yield() >>> []
               end
-
-              SCF.yield() >>> []
             end
           end
-        end >>> []
+        end >>> ret_t
       end
 
     {v, state, env}

lib/charms/prelude.ex

@@ -8,7 +8,8 @@ defmodule Charms.Prelude do
     @enif_functions ++ [:result_at] ++ @binary_ops
   end
 
-  defp constant_of_same_type(i, v, opts) do
+  @doc false
+  def constant_of_same_type(i, v, opts) do
     mlir ctx: opts[:ctx], block: opts[:block] do
       t = MLIR.CAPI.mlirValueGetType(v)
       Arith.constant(value: Attribute.integer(t, i)) >>> t

test/if_test.exs

@@ -0,0 +1,32 @@
+defmodule IfTest do
+  use ExUnit.Case
+
+  test "if with value" do
+    defmodule GetIntIf do
+      use Charms
+      alias Charms.{Pointer, Term}
+
+      defm get(env, i) :: Term.t() do
+        zero = arith.constant(value: Attribute.integer(i32(), 0))
+        one = arith.constant(value: Attribute.integer(i32(), 1))
+        i_ptr = Pointer.allocate(i32())
+        enif_get_int(env, i, i_ptr)
+        i = Pointer.load(i32(), i_ptr)
+
+        ret =
+          if(i > 0) do
+            one
+          else
+            zero
+          end
+
+        ret = enif_make_int(env, ret)
+        func.return(ret)
+      end
+    end
+    |> Charms.JIT.init()
+
+    assert GetIntIf.get(100) == 1
+    assert GetIntIf.get(-100) == 0
+  end
+end