diff --git a/src/Lean/Compiler/ExternAttr.lean b/src/Lean/Compiler/ExternAttr.lean index 440d5b812c..44e9278ce7 100644 --- a/src/Lean/Compiler/ExternAttr.lean +++ b/src/Lean/Compiler/ExternAttr.lean @@ -7,6 +7,7 @@ import Lean.Expr import Lean.Environment import Lean.Attributes import Lean.ProjFns +import Lean.Meta.Basic namespace Lean @@ -28,10 +29,10 @@ inductive ExternEntry - `@[extern cpp "foo" llvm adhoc]` encoding: ```.entries = [standard `cpp "foo", adhoc `llvm]``` - `@[extern 2 cpp "io_prim_println"]` - encoding: ```.arity = 2, .entries = [standard `cpp "ioPrimPrintln"]``` + encoding: ```.arity? = 2, .entries = [standard `cpp "ioPrimPrintln"]``` -/ structure ExternAttrData := -(arity : Option Nat := none) +(arity? : Option Nat := none) (entries : List ExternEntry) instance ExternAttrData.inhabited : Inhabited ExternAttrData := ⟨{ entries := [] }⟩ @@ -67,11 +68,11 @@ match s with match (args.get! i).isStrLit? with | some str => if args.size == i+1 then - Except.ok { arity := arity, entries := [ ExternEntry.standard `all str ] } + Except.ok { arity? := arity, entries := [ ExternEntry.standard `all str ] } else Except.error "invalid extern attribute" | none => match syntaxToExternEntries args i [] with - | Except.ok entries => Except.ok { arity := arity, entries := entries } + | Except.ok entries => Except.ok { arity? := arity, entries := entries } | Except.error msg => Except.error msg | _ => Except.error "unexpected kind of argument" @@ -158,4 +159,21 @@ match entry with | ExternEntry.foreign _ n => pure n | _ => failure +def getExternConstArity (declName : Name) : CoreM (Option Nat) := do +env ← getEnv; +match getExternAttrData env declName with +| none => pure none +| some data => match data.arity? with + | some arity => pure arity + | none => do + cinfo ← getConstInfo declName; + (arity, _) ← (Meta.forallTelescopeReducing cinfo.type fun xs _ => pure xs.size : MetaM Nat).run; + pure (some arity) + +@[export lean_get_extern_const_arity] +def getExternConstArityExport (env : Environment) (declName : Name) : IO (Option Nat) := +catch + (do (arity?, _) ← (getExternConstArity declName).toIO {} { env := env }; pure arity?) + (fun _ => pure none) + end Lean diff --git a/src/library/compiler/extern_attribute.cpp b/src/library/compiler/extern_attribute.cpp index f73251e2e8..3d3b6ce52f 100644 --- a/src/library/compiler/extern_attribute.cpp +++ b/src/library/compiler/extern_attribute.cpp @@ -7,6 +7,8 @@ Authors: Leonardo de Moura #include #include #include "util/object_ref.h" +#include "util/option_ref.h" +#include "util/io.h" #include "kernel/instantiate.h" #include "kernel/type_checker.h" #include "library/util.h" @@ -27,52 +29,15 @@ bool is_extern_constant(environment const & env, name const & c) { return static_cast(get_extern_attr_data(env, c)); } -static optional get_given_arity(environment const & env, name const & c) { - optional v = get_extern_attr_data(env, c); - lean_assert(v); - object * arity = cnstr_get(v->raw(), 0); - if (is_scalar(arity)) return optional(); // none - // arity is (some v) - arity = cnstr_get(arity, 0); - if (is_scalar(arity)) return optional(unbox(arity)); - return optional(); // ignore big nums -} - -bool is_io(expr const & type) { - return is_app_of(type, get_io_name()) || is_app_of(type, get_eio_name()) || is_app_of(type, get_st_name()); -} - -/* - Similar to lean::get_arity, but adds `1` if resultant type is of the form `IO a`. - Motivation: we want to correctly guess the IO extern primitive without user help. - By user help, we mean a definition such as - ``` - @[extern 2 "lean_io_getenv"] - constant getEnv (var : @& String) : IO (Option String) := default _ - ``` - were we explicitly say `getEnv` has arity 2. -*/ -static unsigned get_arity_for_extern(expr type) { - unsigned r = 0; - while (is_pi(type)) { - type = binding_body(type); - r++; - } - if (is_io(type)) { - r++; - } - return r; -} +extern "C" object * lean_get_extern_const_arity(object* env, object*, object* w); optional get_extern_constant_arity(environment const & env, name const & c) { - if (is_extern_constant(env, c)) { - if (optional given_arity = get_given_arity(env, c)) { - return given_arity; - } - /* Infer arity from type */ - return optional(get_arity_for_extern(env.get(c).get_type())); + auto arity = get_io_result>(lean_get_extern_const_arity(env.to_obj_arg(), c.to_obj_arg(), lean_io_mk_world())); + if (optional aux = arity.get()) { + return optional(aux->get_small_value()); + } else { + return optional(); } - return optional(); } bool get_extern_borrowed_info(environment const & env, name const & c, buffer & borrowed_args, bool & borrowed_res) { @@ -87,18 +52,14 @@ bool get_extern_borrowed_info(environment const & env, name const & c, buffer given_arity = get_given_arity(env, c)) { - if (*given_arity < arity) { - borrowed_args.shrink(*given_arity); + if (optional c_arity = get_extern_constant_arity(env, c)) { + if (*c_arity < arity) { + borrowed_args.shrink(*c_arity); return true; - } else if (*given_arity > arity) { - borrowed_args.resize(*given_arity, false); + } else if (*c_arity > arity) { + borrowed_args.resize(*c_arity, false); return true; } - } else if (is_io(type)) { - /* See: `get_arity_for_extern`. We have special code for guessing - the arity for external IO primitives. */ - borrowed_args.push_back(false); } borrowed_res = is_borrowed(type); return true; @@ -126,23 +87,17 @@ optional get_extern_constant_ll_type(environment const & env, name const & } type = instantiate_rev(type, locals.size(), locals.data()); expr ll_type; - if (optional given_arity = get_given_arity(env, c)) { - if (arity < *given_arity) { + if (optional c_arity = get_extern_constant_arity(env, c)) { + if (arity < *c_arity) { /* Fill with `_obj` */ - arg_ll_types.resize(*given_arity, mk_enf_object_type()); + arg_ll_types.resize(*c_arity, mk_enf_object_type()); ll_type = mk_enf_object_type(); - } else if (arity > *given_arity) { - arg_ll_types.shrink(*given_arity); + } else if (arity > *c_arity) { + arg_ll_types.shrink(*c_arity); ll_type = mk_enf_object_type(); /* Result is a closure */ } else { ll_type = mk_runtime_type(st, lctx, type); } - } else if (is_io(type)) { - /* Add "world". - See: `get_arity_for_extern`. We have special code for guessing - the arity for external IO primitives. */ - arg_ll_types.push_back(mk_enf_object_type()); - ll_type = mk_enf_object_type(); } else { ll_type = mk_runtime_type(st, lctx, type); }