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

@@ -570,6 +570,29 @@ defmodule Charms.Defm.Expander do
     Module.concat(mod, func)
   end
 
+  defp expand_if_clause_body(nil, state, _env) do
+    mlir ctx: state.mlir.ctx, block: state.mlir.blk do
+      SCF.yield() >>> []
+      []
+    end
+  end
+
+  defp expand_if_clause_body(clause_body, state, env) do
+    mlir ctx: state.mlir.ctx, block: state.mlir.blk do
+      {ret, _, _} = expand(clause_body, state, env)
+
+      case ret |> List.wrap() |> List.last() do
+        %MLIR.Operation{} ->
+          SCF.yield() >>> []
+          []
+
+        %MLIR.Value{} = last ->
+          SCF.yield(last) >>> []
+          MLIR.Value.type(last)
+      end
+    end
+  end
+
   ## Macro handling
 
   # This is going to be the function where you will intercept expansions
@@ -711,7 +734,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 +743,24 @@ 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 =
+              expand_if_clause_body(true_body, put_in(state.mlir.blk, Beaver.Env.block()), env)
+          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)
-              end
-
-              SCF.yield() >>> []
+              expand_if_clause_body(false_body, put_in(state.mlir.blk, Beaver.Env.block()), env)
             end
           end
-        end >>> []
+        end >>> ret_t
       end
 
     {v, state, env}

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