Skip to main content

aion/runtime/
nif.rs

1//! NIF registration surface.
2
3use std::collections::BTreeSet;
4
5use beamr::native::NativeFn;
6
7/// Module/function/arity key for a native implemented function.
8#[derive(Clone, Debug, Eq, Hash, PartialEq)]
9pub struct Mfa {
10    /// BEAM module name that owns the native function import.
11    pub module: String,
12    /// BEAM function name imported from `module`.
13    pub function: String,
14    /// Function arity.
15    pub arity: u8,
16}
17
18impl Mfa {
19    /// Construct an MFA key from string-like module and function names.
20    #[must_use]
21    pub fn new(module: impl Into<String>, function: impl Into<String>, arity: u8) -> Self {
22        Self {
23            module: module.into(),
24            function: function.into(),
25            arity,
26        }
27    }
28
29    /// Return the human-readable MFA as `module:function/arity`.
30    #[must_use]
31    pub fn display(&self) -> String {
32        format!("{}:{}/{}", self.module, self.function, self.arity)
33    }
34}
35
36/// A host- or engine-owned native implemented function entry.
37#[derive(Clone, Debug)]
38pub struct NifEntry {
39    /// MFA key used by BEAM import resolution.
40    pub mfa: Mfa,
41    /// Rust function pointer compatible with beamr's native registry.
42    pub function: NativeFn,
43    /// Whether beamr should mark the entry for dirty scheduler execution.
44    pub is_dirty: bool,
45}
46
47impl NifEntry {
48    /// Construct a normal native implemented function entry.
49    #[must_use]
50    pub fn new(mfa: Mfa, function: NativeFn) -> Self {
51        Self {
52            mfa,
53            function,
54            is_dirty: false,
55        }
56    }
57
58    /// Construct a native implemented function entry marked dirty.
59    #[must_use]
60    pub fn dirty(mfa: Mfa, function: NativeFn) -> Self {
61        Self {
62            mfa,
63            function,
64            is_dirty: true,
65        }
66    }
67}
68
69/// Accumulates NIF entries before they are installed into the runtime.
70#[derive(Clone, Debug, Default)]
71pub struct NifRegistration {
72    entries: Vec<NifEntry>,
73}
74
75#[cfg(test)]
76pub(crate) fn test_native_zero(
77    args: &[beamr::term::Term],
78    context: &mut beamr::native::ProcessContext,
79) -> Result<beamr::term::Term, beamr::term::Term> {
80    let _ = context;
81    if args.len() > 255 {
82        return Err(beamr::term::Term::small_int(0));
83    }
84    Ok(beamr::term::Term::small_int(0))
85}
86
87impl NifRegistration {
88    /// Construct an empty NIF registration collection.
89    #[must_use]
90    pub const fn new() -> Self {
91        Self {
92            entries: Vec::new(),
93        }
94    }
95
96    /// Add host-supplied NIF entries to the collection.
97    pub fn add_host_nifs(&mut self, entries: impl IntoIterator<Item = NifEntry>) -> &mut Self {
98        self.entries.extend(entries);
99        self
100    }
101
102    /// Add engine-owned NIF entries to the collection.
103    ///
104    /// Registers the `aion_flow_ffi` NIFs that back the Gleam SDK's
105    /// `@external(erlang, "aion_flow_ffi", ...)` declarations.
106    pub fn add_engine_nifs(&mut self) -> &mut Self {
107        self.entries
108            .extend(super::engine_nifs::engine_nif_entries());
109        self
110    }
111
112    /// Return the unique module names represented by the collected NIF entries.
113    ///
114    /// These names are derived from each entry's MFA and are therefore kept in
115    /// sync with both engine-owned and host-supplied registrations.
116    #[must_use]
117    pub fn module_names(&self) -> Vec<String> {
118        self.entries
119            .iter()
120            .map(|entry| entry.mfa.module.clone())
121            .collect::<BTreeSet<_>>()
122            .into_iter()
123            .collect()
124    }
125
126    /// Consume the collection and return the entries to install.
127    #[must_use]
128    pub fn into_entries(self) -> Vec<NifEntry> {
129        self.entries
130    }
131
132    /// Return true when no NIF entries have been collected.
133    #[must_use]
134    pub fn is_empty(&self) -> bool {
135        self.entries.is_empty()
136    }
137
138    /// Return the number of NIF entries collected.
139    #[must_use]
140    pub fn len(&self) -> usize {
141        self.entries.len()
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use beamr::native::ProcessContext;
148    use beamr::term::Term;
149
150    use super::{Mfa, NifEntry, NifRegistration};
151
152    fn native_zero(args: &[Term], context: &mut ProcessContext) -> Result<Term, Term> {
153        super::test_native_zero(args, context)
154    }
155
156    #[test]
157    fn collects_host_and_engine_nifs() {
158        let mut registration = NifRegistration::new();
159        registration
160            .add_engine_nifs()
161            .add_host_nifs([NifEntry::dirty(Mfa::new("host", "zero", 0), native_zero)]);
162
163        assert!(registration.len() >= 2);
164        let module_names = registration.module_names();
165        assert!(module_names.iter().any(|module| module == "aion_flow_ffi"));
166        assert!(module_names.iter().any(|module| module == "host"));
167
168        let entries = registration.into_entries();
169        let host_entry = entries.iter().find(|e| e.mfa.display() == "host:zero/0");
170        assert!(host_entry.is_some_and(|e| e.is_dirty));
171        let dispatch_activity = entries
172            .iter()
173            .find(|e| e.mfa.display() == "aion_flow_ffi:dispatch_activity/3");
174        assert!(dispatch_activity.is_some_and(|e| !e.is_dirty));
175        let await_activity = entries
176            .iter()
177            .find(|e| e.mfa.display() == "aion_flow_ffi:await_activity_result/1");
178        assert!(await_activity.is_some_and(|e| !e.is_dirty));
179    }
180}