ensure enif function has correct signature

beaver-lodge · Oct 8, 2024 · ceff98e · ceff98e
1 parent af03d08
commit ceff98e
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 11 deletions.
diff --git a/lib/charms/defm/expander.ex b/lib/charms/defm/expander.ex
@@ -402,10 +402,12 @@ defmodule Charms.Defm.Expander do
   defp expand({fun, _meta, [left, right]}, state, env) when fun in @intrinsics do
     {left, state, env} = expand(left, state, env)
     {right, state, env} = expand(right, state, env)
+    loc = MLIR.Location.from_env(env)
 
     {Charms.Prelude.handle_intrinsic(fun, [left, right],
        ctx: state.mlir.ctx,
-       block: state.mlir.blk
+       block: state.mlir.blk,
+       loc: loc
      ), state, env}
   end
 
@@ -463,6 +465,7 @@ defmodule Charms.Defm.Expander do
     arity = length(args)
     mfa = {module, fun, arity}
     state = update_in(state.remotes, &[mfa | &1])
+    loc = MLIR.Location.from_env(env)
 
     if is_atom(module) do
       try do
@@ -480,8 +483,11 @@ defmodule Charms.Defm.Expander do
               function_exported?(module, :__intrinsics__, 0) and fun in module.__intrinsics__() ->
                 {args, state, env} = expand(args, state, env)
 
-                {module.handle_intrinsic(fun, args, ctx: state.mlir.ctx, block: state.mlir.blk),
-                 state, env}
+                {module.handle_intrinsic(fun, args,
+                   ctx: state.mlir.ctx,
+                   block: state.mlir.blk,
+                   loc: loc
+                 ), state, env}
 
               module == MLIR.Attribute ->
                 {args, state, env} = expand(args, state, env)
@@ -499,6 +505,7 @@ defmodule Charms.Defm.Expander do
                   attr = unquote(attr)
                   term_ptr = Pointer.allocate(Term.t())
                   size = String.length(attr)
+                  size = value index.casts(size) :: i64()
                   buffer_ptr = Pointer.allocate(i8(), size)
                   buffer = ptr_to_memref(buffer_ptr, size)
                   memref.copy(attr, buffer)
@@ -1057,11 +1064,13 @@ defmodule Charms.Defm.Expander do
   ## Helpers
 
   defp expand_remote(_meta, Kernel, fun, args, state, env) when fun in @intrinsics do
+    loc = MLIR.Location.from_env(env)
     {args, state, env} = expand(args, state, env)
 
     {Charms.Prelude.handle_intrinsic(fun, args,
        ctx: state.mlir.ctx,
-       block: state.mlir.blk
+       block: state.mlir.blk,
+       loc: loc
      ), state, env}
   end
 
@@ -1087,14 +1096,16 @@ defmodule Charms.Defm.Expander do
     state = update_in(state.locals, &[{fun, length(args)} | &1])
     {args, state, env} = expand_list(args, state, env)
     Code.ensure_loaded!(MLIR.Type)
+    loc = MLIR.Location.from_env(env)
 
     if function_exported?(MLIR.Type, fun, 1) do
       {apply(MLIR.Type, fun, [[ctx: state.mlir.ctx]]), state, env}
     else
       case i =
              Charms.Prelude.handle_intrinsic(fun, args,
                ctx: state.mlir.ctx,
-               block: state.mlir.blk
+               block: state.mlir.blk,
+               loc: loc
              ) do
         :not_handled ->
           create_call(env.module, fun, args, [], state, env)

diff --git a/lib/charms/defm/pass/create_absent_func.ex b/lib/charms/defm/pass/create_absent_func.ex
@@ -42,10 +42,18 @@ defmodule Charms.Defm.Pass.CreateAbsentFunc do
                name_str <- MLIR.StringRef.to_string(name),
                false <- MapSet.member?(created, name_str) do
             mlir ctx: ctx, block: block do
+              ft =
+                if s = Beaver.ENIF.signature(ctx, String.to_atom(name_str)) do
+                  {arg_types, ret_types} = s
+                  Type.function(arg_types, ret_types)
+                else
+                  Type.function(arg_types, ret_types)
+                end
+
               Func.func _(
-                          sym_name: "\"#{name_str}\"",
+                          sym_name: MLIR.Attribute.string(name_str),
                           sym_visibility: MLIR.Attribute.string("private"),
-                          function_type: Type.function(arg_types, ret_types)
+                          function_type: ft
                         ) do
                 region do
                 end

diff --git a/lib/charms/intrinsic.ex b/lib/charms/intrinsic.ex
@@ -3,7 +3,7 @@ defmodule Charms.Intrinsic do
   Behaviour to define intrinsic functions.
   """
   alias Beaver
-  @type opt :: {:ctx, MLIR.Context.t()} | {:block, MLIR.Block.t()}
+  @type opt :: {:ctx, MLIR.Context.t()} | {:block, MLIR.Block.t() | {:loc, MLIR.Location.t()}}
   @type opts :: [opt | {atom(), term()}]
   @type ir_return :: MLIR.Value.t() | MLIR.Operation.t()
   @type intrinsic_return :: ir_return() | (any() -> ir_return())

diff --git a/lib/charms/prelude.ex b/lib/charms/prelude.ex
@@ -92,7 +92,7 @@ defmodule Charms.Prelude do
     args = args |> Enum.zip(arg_types) |> Enum.map(&wrap_arg(&1, opts))
 
     mlir ctx: opts[:ctx], block: opts[:block] do
-      Func.call(args, callee: Attribute.flat_symbol_ref("#{name}")) >>>
+      Func.call(args, callee: Attribute.flat_symbol_ref("#{name}"), loc: opts[:loc]) >>>
         case ret_types do
           [ret] ->
             ret

diff --git a/test/string_test.exs b/test/string_test.exs
@@ -10,8 +10,9 @@ defmodule StringTest do
         str = "this is a string"
         str = "this is a string"
         term_ptr = Pointer.allocate(Term.t())
-        d_ptr = enif_make_new_binary(env, String.length(str), term_ptr)
-        m = ptr_to_memref(d_ptr)
+        size = value index.casts(String.length(str)) :: i64()
+        d_ptr = enif_make_new_binary(env, size, term_ptr)
+        m = ptr_to_memref(d_ptr, size)
         memref.copy(str, m)
         t = Pointer.load(Term.t(), term_ptr)
         func.return(t)