Skip to main content

sim_lib_numbers_func/implementation/
function.rs

1//! Function operations: the `fn`, `call`, and `grad` callables and the
2//! function class builder backing the `Func` domain.
3
4use std::{any::Any, sync::Arc};
5
6use sim_kernel::{
7    Args, Callable, Class, ClassId, ClassRef, Cx, DefaultFactory, Env, Error, Expr, Factory,
8    Object, ObjectEncode, ObjectEncoding, ReadConstructor, ReadConstructorRef, Result, ShapeRef,
9    Symbol, TableRef, Value,
10};
11use sim_lib_numbers_cas::{cas_expr_to_surface_expr, expr_to_cas_expr};
12
13use super::domain::{func_class_symbol, value_shape_symbol};
14use super::value::{Func, build_func_value};
15
16/// Returns the symbol bound to the `fn` function-builder callable.
17///
18/// # Examples
19///
20/// ```
21/// use sim_lib_numbers_func::{call_symbol, fn_symbol, grad_symbol};
22///
23/// assert_eq!(fn_symbol().to_string(), "fn");
24/// assert_eq!(call_symbol().to_string(), "call");
25/// assert_eq!(grad_symbol().to_string(), "grad");
26/// ```
27pub fn fn_symbol() -> Symbol {
28    Symbol::new("fn")
29}
30
31/// Returns the symbol bound to the `call` apply-a-function callable.
32pub fn call_symbol() -> Symbol {
33    Symbol::new("call")
34}
35
36/// Returns the symbol bound to the `grad` gradient-of-a-function callable.
37pub fn grad_symbol() -> Symbol {
38    Symbol::new("grad")
39}
40
41#[derive(Clone)]
42pub struct FnBuilder;
43
44impl Object for FnBuilder {
45    fn display(&self, _cx: &mut Cx) -> Result<String> {
46        Ok("#<function fn>".to_owned())
47    }
48
49    fn as_any(&self) -> &dyn Any {
50        self
51    }
52}
53
54impl sim_kernel::ObjectCompat for FnBuilder {
55    fn class(&self, cx: &mut Cx) -> Result<ClassRef> {
56        function_class(cx)
57    }
58    fn as_expr(&self, _cx: &mut Cx) -> Result<Expr> {
59        Ok(Expr::Symbol(fn_symbol()))
60    }
61    fn as_callable(&self) -> Option<&dyn Callable> {
62        Some(self)
63    }
64}
65
66impl Callable for FnBuilder {
67    fn call(&self, _cx: &mut Cx, _args: Args) -> Result<Value> {
68        Err(Error::Eval(
69            "fn must be called with unevaluated parameters and a body".to_owned(),
70        ))
71    }
72
73    fn call_exprs(&self, cx: &mut Cx, args: sim_kernel::RawArgs) -> Result<Value> {
74        let args = args.into_exprs();
75        let [vars_expr, body_expr] = args.as_slice() else {
76            return Err(Error::Eval(
77                "fn expects exactly a parameter list and one body expression".to_owned(),
78            ));
79        };
80        let vars = parse_vars_expr(vars_expr)?;
81        let body_cas = expr_to_cas_expr(cx, body_expr)?
82            .ok_or_else(|| Error::Eval("fn body must be CAS-compatible".to_owned()))?;
83        build_func_value(cx, Func::symbolic(vars, body_cas))
84    }
85}
86
87#[derive(Clone)]
88pub struct CallFunction;
89
90impl Object for CallFunction {
91    fn display(&self, _cx: &mut Cx) -> Result<String> {
92        Ok("#<function call>".to_owned())
93    }
94
95    fn as_any(&self) -> &dyn Any {
96        self
97    }
98}
99
100impl sim_kernel::ObjectCompat for CallFunction {
101    fn class(&self, cx: &mut Cx) -> Result<ClassRef> {
102        function_class(cx)
103    }
104    fn as_expr(&self, _cx: &mut Cx) -> Result<Expr> {
105        Ok(Expr::Symbol(call_symbol()))
106    }
107    fn as_callable(&self) -> Option<&dyn Callable> {
108        Some(self)
109    }
110}
111
112impl Callable for CallFunction {
113    fn call(&self, cx: &mut Cx, args: Args) -> Result<Value> {
114        let mut values = args.into_vec();
115        if values.is_empty() {
116            return Err(Error::Eval(
117                "call expects a callable value and at least zero arguments".to_owned(),
118            ));
119        }
120        let callable = values.remove(0);
121        cx.call_value(callable, Args::new(values))
122    }
123}
124
125#[derive(Clone)]
126pub struct GradFunction;
127
128impl Object for GradFunction {
129    fn display(&self, _cx: &mut Cx) -> Result<String> {
130        Ok("#<function grad>".to_owned())
131    }
132
133    fn as_any(&self) -> &dyn Any {
134        self
135    }
136}
137
138impl sim_kernel::ObjectCompat for GradFunction {
139    fn class(&self, cx: &mut Cx) -> Result<ClassRef> {
140        function_class(cx)
141    }
142    fn as_expr(&self, _cx: &mut Cx) -> Result<Expr> {
143        Ok(Expr::Symbol(grad_symbol()))
144    }
145    fn as_callable(&self) -> Option<&dyn Callable> {
146        Some(self)
147    }
148}
149
150impl Callable for GradFunction {
151    fn call(&self, cx: &mut Cx, args: Args) -> Result<Value> {
152        let values = args.into_vec();
153        let [value] = values.as_slice() else {
154            return Err(Error::Eval(
155                "grad expects exactly one function value".to_owned(),
156            ));
157        };
158        let func = expect_func(value)?;
159        let mut grads = Vec::with_capacity(func.vars.len());
160        for var in &func.vars {
161            let var_value = cx.factory().symbol(var.clone())?;
162            grads.push(cx.call_function(
163                &Symbol::new("diff"),
164                Args::new(vec![value.clone(), var_value]),
165            )?);
166        }
167        cx.factory().list(grads)
168    }
169}
170
171pub(crate) struct FuncValueClass {
172    id: std::sync::atomic::AtomicU32,
173}
174
175pub(crate) fn build_func_class() -> Arc<FuncValueClass> {
176    Arc::new(FuncValueClass {
177        id: std::sync::atomic::AtomicU32::new(0),
178    })
179}
180
181impl FuncValueClass {
182    pub(crate) fn set_id(&self, id: ClassId) {
183        self.id.store(id.0, std::sync::atomic::Ordering::Relaxed);
184    }
185}
186
187impl Object for FuncValueClass {
188    fn display(&self, _cx: &mut Cx) -> Result<String> {
189        Ok(format!("#<class {}>", func_class_symbol()))
190    }
191
192    fn as_any(&self) -> &dyn Any {
193        self
194    }
195}
196
197impl sim_kernel::ObjectCompat for FuncValueClass {
198    fn class(&self, cx: &mut Cx) -> Result<ClassRef> {
199        if let Some(value) = cx
200            .registry()
201            .class_by_symbol(&Symbol::qualified("core", "Class"))
202        {
203            return Ok(value.clone());
204        }
205        DefaultFactory.class_stub(
206            sim_kernel::CORE_CLASS_CLASS_ID,
207            Symbol::qualified("core", "Class"),
208        )
209    }
210    fn as_expr(&self, _cx: &mut Cx) -> Result<Expr> {
211        Ok(Expr::Symbol(func_class_symbol()))
212    }
213    fn as_callable(&self) -> Option<&dyn Callable> {
214        Some(self)
215    }
216    fn as_class(&self) -> Option<&dyn Class> {
217        Some(self)
218    }
219    fn as_read_constructor(&self) -> Option<&dyn ReadConstructor> {
220        Some(self)
221    }
222}
223
224impl Callable for FuncValueClass {
225    fn call(&self, cx: &mut Cx, args: Args) -> Result<Value> {
226        let values = args.into_vec();
227        let [vars_value, body_value] = values.as_slice() else {
228            return Err(Error::Eval(format!(
229                "class {} expects exactly two arguments",
230                func_class_symbol()
231            )));
232        };
233        let vars = parse_vars_value(cx, vars_value)?;
234        let body_expr = body_value.object().as_expr(cx)?;
235        let body_cas = expr_to_cas_expr(cx, &body_expr)?
236            .ok_or_else(|| Error::Eval("numbers/Func body must be CAS-compatible".to_owned()))?;
237        build_func_value(cx, Func::symbolic(vars, body_cas))
238    }
239}
240
241impl Class for FuncValueClass {
242    fn id(&self) -> ClassId {
243        ClassId(self.id.load(std::sync::atomic::Ordering::Relaxed))
244    }
245
246    fn symbol(&self) -> Symbol {
247        func_class_symbol()
248    }
249
250    fn constructor_shape(&self, cx: &mut Cx) -> Result<ShapeRef> {
251        cx.factory().nil()
252    }
253
254    fn instance_shape(&self, cx: &mut Cx) -> Result<ShapeRef> {
255        Ok(cx
256            .registry()
257            .shape_by_symbol(&value_shape_symbol())
258            .cloned()
259            .unwrap_or(cx.factory().symbol(value_shape_symbol())?))
260    }
261
262    fn read_constructor(&self, cx: &mut Cx) -> Result<Option<ReadConstructorRef>> {
263        Ok(cx.registry().class_by_symbol(&func_class_symbol()).cloned())
264    }
265
266    fn members(&self, cx: &mut Cx) -> Result<TableRef> {
267        cx.factory().table(Vec::new())
268    }
269}
270
271impl ReadConstructor for FuncValueClass {
272    fn symbol(&self) -> Symbol {
273        func_class_symbol()
274    }
275
276    fn args_shape(&self, cx: &mut Cx) -> Result<ShapeRef> {
277        cx.factory().nil()
278    }
279
280    fn construct_read(&self, cx: &mut Cx, args: Vec<Value>) -> Result<Value> {
281        self.call(cx, Args::new(args))
282    }
283}
284
285impl ObjectEncode for Func {
286    fn object_encoding(&self, cx: &mut Cx) -> Result<ObjectEncoding> {
287        let Some(body_cas) = self.body_cas() else {
288            return Err(Error::Eval(
289                "native-only functions do not have a read-construct encoding".to_owned(),
290            ));
291        };
292        Ok(ObjectEncoding::Constructor {
293            class: func_class_symbol(),
294            args: vec![
295                vars_expr(&self.vars),
296                cas_expr_to_surface_expr(cx, body_cas)?,
297            ],
298        })
299    }
300}
301
302pub(crate) fn parse_vars_expr(expr: &Expr) -> Result<Vec<Symbol>> {
303    let Expr::List(items) = expr else {
304        return Err(Error::Eval(
305            "function parameter list must be a list of symbols".to_owned(),
306        ));
307    };
308    items
309        .iter()
310        .map(|item| match item {
311            Expr::Symbol(symbol) => Ok(symbol.clone()),
312            _ => Err(Error::Eval(
313                "function parameter list must contain only symbols".to_owned(),
314            )),
315        })
316        .collect()
317}
318
319fn parse_vars_value(cx: &mut Cx, value: &Value) -> Result<Vec<Symbol>> {
320    parse_vars_expr(&value.object().as_expr(cx)?)
321}
322
323pub(crate) fn vars_expr(vars: &[Symbol]) -> Expr {
324    Expr::List(vars.iter().cloned().map(Expr::Symbol).collect())
325}
326
327pub(crate) fn function_class(cx: &mut Cx) -> Result<ClassRef> {
328    if let Some(value) = cx
329        .registry()
330        .class_by_symbol(&Symbol::qualified("core", "Function"))
331    {
332        return Ok(value.clone());
333    }
334    cx.factory().class_stub(
335        sim_kernel::CORE_FUNCTION_CLASS_ID,
336        Symbol::qualified("core", "Function"),
337    )
338}
339
340pub(crate) fn expect_func(value: &Value) -> Result<&Func> {
341    value
342        .object()
343        .downcast_ref::<Func>()
344        .ok_or_else(|| Error::Eval("expected a numbers/func value".to_owned()))
345}
346
347pub(crate) fn child_env_with_args(parent: &Env, vars: &[Symbol], args: &[Value]) -> Result<Env> {
348    if vars.len() != args.len() {
349        return Err(Error::Eval(format!(
350            "function expected {} arguments but received {}",
351            vars.len(),
352            args.len()
353        )));
354    }
355    let mut env = Env::child(Arc::new(parent.clone()));
356    for (var, value) in vars.iter().cloned().zip(args.iter().cloned()) {
357        env.define(var, value);
358    }
359    Ok(env)
360}