pounce_nl/
nl_external.rs

1//! AMPL imported (external) function support via the `funcadd_ASL` ABI.
2//!
3//! This module implements enough of AMPL's `funcadd.h` ABI to:
4//!
5//! 1. `dlopen` a user-supplied shared library;
6//! 2. resolve the `funcadd_ASL` symbol and call it;
7//! 3. receive registration callbacks of the form `Addfunc(name, rfunc, type,
8//!    nargs, funcinfo, ae)` and record them;
9//! 4. later call back into the registered `rfunc` with an `arglist` to obtain
10//!    function values, gradients, and Hessians.
11//!
12//! The `AmplExports` and `Arglist` struct layouts are taken from
13//! AMPL-MP/ASL `funcadd.h`; cross-checked against the ctypes mapping in
14//! `pyomo.core.base.external`. Fields we don't populate are left null —
15//! Pyomo does the same and it is sufficient for IDAES's Helmholtz library
16//! (see issue #15).
17//!
18//! All unsafe FFI is contained in this module. Public surface is safe.
19
20use std::collections::HashMap;
21use std::ffi::{c_char, c_int, c_long, c_void, CStr, CString};
22use std::path::Path;
23use std::ptr;
24use std::sync::{Arc, Mutex, OnceLock};
25
26use libloading::{Library, Symbol};
27
28use crate::nl_reader::{Expr, FuncallArg, ImportedFunc};
29
30/// Resolved AMPL imported function: shared library + registered name.
31/// `NlProblem` carries one of these per `ImportedFunc` id when external
32/// functions are wired up at problem-build time. The same `Arc<ExternalLibrary>`
33/// may be shared across many funcall ids (one library typically registers
34/// several functions).
35#[derive(Default, Clone)]
36pub struct ExternalResolver {
37    /// `Funcall { id }` -> (library, registered function name).
38    pub funcs_by_id: HashMap<usize, (Arc<ExternalLibrary>, String)>,
39}
40
41impl ExternalResolver {
42    pub fn is_empty(&self) -> bool {
43        self.funcs_by_id.is_empty()
44    }
45
46    /// Build a resolver for every `ImportedFunc` declared in the `.nl` file
47    /// that is *actually referenced* somewhere in the problem's expressions.
48    ///
49    /// Library paths are resolved through the `AMPLFUNC` environment variable
50    /// (a `\n`-separated list of shared-library paths, matching AMPL/IPOPT
51    /// conventions). Each path is loaded once and queried for every name we
52    /// need. Returns an error if a referenced name cannot be found in any
53    /// listed library, or if `AMPLFUNC` is missing.
54    pub fn build_for_problem(
55        imported_funcs: &[ImportedFunc],
56        referenced_ids: &std::collections::BTreeSet<usize>,
57    ) -> Result<Self, String> {
58        if referenced_ids.is_empty() {
59            return Ok(Self::default());
60        }
61        let amplfunc = std::env::var("AMPLFUNC").map_err(|_| {
62            "problem uses external functions but AMPLFUNC is not set; \
63             set AMPLFUNC to a newline-separated list of AMPL shared-library paths"
64                .to_string()
65        })?;
66        let mut libs: Vec<Arc<ExternalLibrary>> = Vec::new();
67        for path_str in amplfunc
68            .split('\n')
69            .map(|s| s.trim())
70            .filter(|s| !s.is_empty())
71        {
72            let path = std::path::Path::new(path_str);
73            let lib = ExternalLibrary::load(path).map_err(|e| format!("AMPLFUNC: {e}"))?;
74            libs.push(Arc::new(lib));
75        }
76
77        let mut funcs_by_id: HashMap<usize, (Arc<ExternalLibrary>, String)> = HashMap::new();
78        for id in referenced_ids {
79            let decl = imported_funcs
80                .iter()
81                .find(|f| f.id == *id)
82                .ok_or_else(|| format!("funcall id {id} has no F<{id}> declaration"))?;
83            let found = libs
84                .iter()
85                .find(|lib| lib.get(&decl.name).is_some())
86                .ok_or_else(|| {
87                    format!(
88                        "external function '{}' (id {}) not found in any library on AMPLFUNC",
89                        decl.name, decl.id
90                    )
91                })?;
92            funcs_by_id.insert(*id, (found.clone(), decl.name.clone()));
93        }
94        Ok(Self { funcs_by_id })
95    }
96}
97
98/// Walk an `Expr` and collect every funcall id it references (including
99/// through CSEs). Used to build an `ExternalResolver` covering exactly the
100/// functions a problem actually uses.
101pub fn collect_funcall_ids(e: &Expr, out: &mut std::collections::BTreeSet<usize>) {
102    match e {
103        Expr::Const(_) | Expr::Var(_) => {}
104        Expr::Binary(_, a, b) => {
105            collect_funcall_ids(a, out);
106            collect_funcall_ids(b, out);
107        }
108        Expr::Unary(_, a) => collect_funcall_ids(a, out),
109        Expr::Sum(args) | Expr::MinList(args) | Expr::MaxList(args) => {
110            for a in args {
111                collect_funcall_ids(a, out);
112            }
113        }
114        Expr::Compare(_, a, b) | Expr::And(a, b) | Expr::Or(a, b) => {
115            collect_funcall_ids(a, out);
116            collect_funcall_ids(b, out);
117        }
118        Expr::Not(a) => collect_funcall_ids(a, out),
119        Expr::Cond { cond, then_, else_ } => {
120            collect_funcall_ids(cond, out);
121            collect_funcall_ids(then_, out);
122            collect_funcall_ids(else_, out);
123        }
124        Expr::Cse(body) => collect_funcall_ids(body, out),
125        Expr::Funcall { id, args } => {
126            out.insert(*id);
127            for arg in args {
128                if let FuncallArg::Real(e) = arg {
129                    collect_funcall_ids(e, out);
130                }
131            }
132        }
133    }
134}
135
136/// Process-wide lock serialising every call that crosses the AMPL external
137/// ABI. Real AMPL libraries (e.g. IDAES general_helmholtz) keep mutable
138/// global state (cached parameters, tabulated lookups) and are not safe for
139/// concurrent entry. Python's `pyomo.core.base.external` relies on the GIL
140/// for the same guarantee.
141fn ampl_lock() -> &'static Mutex<()> {
142    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
143    LOCK.get_or_init(|| Mutex::new(()))
144}
145
146/// FUNCADD_TYPE bits (mirrors `funcadd.h`).
147pub const FUNCADD_REAL_VALUED: i32 = 0;
148/// Set if the function consumes string arguments. Value is still real.
149pub const FUNCADD_STRING_ARGS: i32 = 1;
150/// Set if the function is allowed to have a variable number of args.
151pub const FUNCADD_OUTPUT_ARGS: i32 = 2;
152pub const FUNCADD_RANDOM_VALUED: i32 = 4;
153
154/// The `arglist` struct from AMPL's `funcadd.h`. Layout must match exactly.
155#[repr(C)]
156pub struct Arglist {
157    pub n: c_int,               // number of args
158    pub nr: c_int,              // number of real input args
159    pub at: *mut c_int,         // argument types
160    pub ra: *mut f64,           // pure real args (IN/OUT/INOUT)
161    pub sa: *mut *const c_char, // symbolic IN args
162    pub derivs: *mut f64,       // partial derivatives (if non-null)
163    pub hes: *mut f64,          // second partials (if non-null)
164    pub dig: *mut c_char,       // skip-derivatives flags
165    pub funcinfo: *mut c_void,  // per-function cookie (set by Addfunc)
166    pub ae: *mut AmplExports,   // points back at our AmplExports
167    pub f: *mut c_void,         // AMPL-internal
168    pub tva: *mut c_void,       // AMPL-internal
169    pub errmsg: *mut c_char,    // error description set by the function
170    pub tmi: *mut c_void,       // Tempmem cookie
171    pub private: *mut c_char,
172    pub nin: c_int,
173    pub nout: c_int,
174    pub nsin: c_int,
175    pub nsout: c_int,
176}
177
178/// Pointer to a user-defined real-valued function, matching
179/// `typedef real (*rfunc)(arglist*)`.
180pub type Rfunc = unsafe extern "C" fn(*mut Arglist) -> f64;
181
182/// Pointer to the `Addfunc` callback provided by the caller.
183pub type AddfuncFn = unsafe extern "C" fn(
184    name: *const c_char,
185    f: Rfunc,
186    ty: c_int,
187    nargs: c_int,
188    funcinfo: *mut c_void,
189    ae: *mut AmplExports,
190);
191
192/// Pointer to the `RandSeedSetter` callback.
193pub type RandSeedSetter = unsafe extern "C" fn(*mut c_void, std::os::raw::c_ulong);
194
195/// Pointer to the `Addrandinit` callback.
196pub type AddrandinitFn =
197    unsafe extern "C" fn(ae: *mut AmplExports, setter: RandSeedSetter, v: *mut c_void);
198
199/// Pointer to the `AtReset` callback.
200pub type AtResetFn = unsafe extern "C" fn(ae: *mut AmplExports, f: *mut c_void, v: *mut c_void);
201
202/// The `AmplExports` struct from AMPL's `funcadd.h`. Layout must match
203/// exactly. Function pointers we don't implement are held as `*mut c_void`
204/// (null) — AMPL's ABI does not require a caller to populate them unless the
205/// loaded library actually invokes them.
206#[repr(C)]
207pub struct AmplExports {
208    pub std_err: *mut c_void,
209    pub addfunc: Option<AddfuncFn>,
210    pub asl_date: c_long,
211    pub fprintf: *mut c_void,
212    pub printf: *mut c_void,
213    pub sprintf: *mut c_void,
214    pub vfprintf: *mut c_void,
215    pub vsprintf: *mut c_void,
216    pub strtod: *mut c_void,
217    pub crypto: *mut c_void,
218    pub asl: *mut c_char,
219    pub at_exit: *mut c_void,
220    pub at_reset: Option<AtResetFn>,
221    pub tempmem: *mut c_void,
222    pub add_table_handler: *mut c_void,
223    pub private_ae: *mut c_char,
224    pub qsortv: *mut c_void,
225
226    pub std_in: *mut c_void,
227    pub std_out: *mut c_void,
228    pub clearerr: *mut c_void,
229    pub fclose: *mut c_void,
230    pub fdopen: *mut c_void,
231    pub feof: *mut c_void,
232    pub ferror: *mut c_void,
233    pub fflush: *mut c_void,
234    pub fgetc: *mut c_void,
235    pub fgets: *mut c_void,
236    pub fileno: *mut c_void,
237    pub fopen: *mut c_void,
238    pub fputc: *mut c_void,
239    pub fputs: *mut c_void,
240    pub fread: *mut c_void,
241    pub freopen: *mut c_void,
242    pub fscanf: *mut c_void,
243    pub fseek: *mut c_void,
244    pub ftell: *mut c_void,
245    pub fwrite: *mut c_void,
246    pub pclose: *mut c_void,
247    pub perror: *mut c_void,
248    pub popen: *mut c_void,
249    pub puts: *mut c_void,
250    pub rewind: *mut c_void,
251    pub scanf: *mut c_void,
252    pub setbuf: *mut c_void,
253    pub setvbuf: *mut c_void,
254    pub sscanf: *mut c_void,
255    pub tempnam: *mut c_void,
256    pub tmpfile: *mut c_void,
257    pub tmpnam: *mut c_void,
258    pub ungetc: *mut c_void,
259    pub ai: *mut c_void,
260    pub getenv: *mut c_void,
261    pub breakfunc: *mut c_void,
262    pub breakarg: *mut c_char,
263
264    // Items available with ASLdate >= 20020501.
265    pub snprintf: *mut c_void,
266    pub vsnprintf: *mut c_void,
267
268    pub addrand: *mut c_void,
269    pub addrandinit: Option<AddrandinitFn>,
270}
271
272// SAFETY: AmplExports itself contains only raw pointers and integers. The
273// library never reads/writes it from another thread concurrently with us
274// (AMPL's model is single-threaded per problem), and we never share it
275// across threads. The Send/Sync bounds only matter because we box the
276// registry inside Arcs.
277unsafe impl Send for AmplExports {}
278unsafe impl Sync for AmplExports {}
279
280/// A function registered by a library via `Addfunc`. Mirrors the ASL
281/// `FUNCADD_TYPE` bits in `funcadd.h`.
282#[derive(Debug, Clone)]
283pub struct RegisteredFunc {
284    pub name: String,
285    pub rfunc: Rfunc,
286    /// OR of FUNCADD_TYPE bits.
287    pub ty: i32,
288    /// Declared arg count. >=0 means exactly that many, <=-1 means "at least
289    /// -(nargs+1) args".
290    pub nargs: i32,
291    /// Cookie set by the library; must be passed through to arglist.funcinfo.
292    pub funcinfo: *mut c_void,
293}
294
295// SAFETY: funcinfo is an opaque cookie owned by the library. We never
296// dereference it; we only pass it back to the library's functions, which
297// expect it. No thread-safety contract is violated by sending the struct.
298unsafe impl Send for RegisteredFunc {}
299unsafe impl Sync for RegisteredFunc {}
300
301impl std::fmt::Debug for ExternalLibrary {
302    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
303        f.debug_struct("ExternalLibrary")
304            .field("funcs", &self.funcs.keys().collect::<Vec<_>>())
305            .finish()
306    }
307}
308
309/// A loaded external-function library plus its registered functions.
310pub struct ExternalLibrary {
311    /// Keep the library alive — it owns the code pages the function pointers
312    /// reference. Arc so `LoadedExternals` can share it.
313    _lib: Arc<Library>,
314    /// The AmplExports we handed to `funcadd_ASL`. Must be kept alive (pinned
315    /// in a Box) because some libraries may capture its address for later
316    /// use (e.g. for `AtReset` bookkeeping).
317    _ae: Box<AmplExports>,
318    /// Registrations collected during `funcadd_ASL`.
319    funcs: HashMap<String, RegisteredFunc>,
320}
321
322impl ExternalLibrary {
323    /// Open a shared library at `path` and invoke its `funcadd_ASL` entry
324    /// point, collecting all functions it registers.
325    pub fn load(path: &Path) -> Result<Self, String> {
326        // Serialise all ABI crossings: library init code and registration
327        // may touch global state that isn't safe under concurrent entry.
328        let _guard = ampl_lock().lock().unwrap_or_else(|e| e.into_inner());
329        // SAFETY: libloading::Library::new is unsafe because it can run
330        // arbitrary initialisers from the shared object. We trust the user's
331        // AMPLFUNC path the same way AMPL/IPOPT do.
332        let lib = unsafe { Library::new(path) }
333            .map_err(|e| format!("failed to open '{}': {}", path.display(), e))?;
334
335        // Resolve `funcadd_ASL`. AMPL's macro `#define funcadd funcadd_ASL`
336        // means every conforming library exports this symbol.
337        type FuncaddFn = unsafe extern "C" fn(*mut AmplExports);
338        let funcadd: Symbol<FuncaddFn> = unsafe { lib.get(b"funcadd_ASL\0") }
339            .map_err(|e| format!("no funcadd_ASL in '{}': {}", path.display(), e))?;
340
341        // Build an AmplExports. Most fields null — the library doesn't call
342        // them (same assumption Pyomo makes). Only the three hooks we can
343        // realistically service are set.
344        let mut ae = Box::new(AmplExports {
345            std_err: ptr::null_mut(),
346            addfunc: Some(trampoline_addfunc),
347            // ASLdate >= 20020501 unlocks the SnprintF/VsnprintF slots.
348            // Pyomo uses 20160307; mirror that.
349            asl_date: 20160307,
350            fprintf: ptr::null_mut(),
351            printf: ptr::null_mut(),
352            sprintf: ptr::null_mut(),
353            vfprintf: ptr::null_mut(),
354            vsprintf: ptr::null_mut(),
355            strtod: ptr::null_mut(),
356            crypto: ptr::null_mut(),
357            asl: ptr::null_mut(),
358            at_exit: ptr::null_mut(),
359            at_reset: Some(trampoline_atreset),
360            tempmem: ptr::null_mut(),
361            add_table_handler: ptr::null_mut(),
362            private_ae: ptr::null_mut(),
363            qsortv: ptr::null_mut(),
364            std_in: ptr::null_mut(),
365            std_out: ptr::null_mut(),
366            clearerr: ptr::null_mut(),
367            fclose: ptr::null_mut(),
368            fdopen: ptr::null_mut(),
369            feof: ptr::null_mut(),
370            ferror: ptr::null_mut(),
371            fflush: ptr::null_mut(),
372            fgetc: ptr::null_mut(),
373            fgets: ptr::null_mut(),
374            fileno: ptr::null_mut(),
375            fopen: ptr::null_mut(),
376            fputc: ptr::null_mut(),
377            fputs: ptr::null_mut(),
378            fread: ptr::null_mut(),
379            freopen: ptr::null_mut(),
380            fscanf: ptr::null_mut(),
381            fseek: ptr::null_mut(),
382            ftell: ptr::null_mut(),
383            fwrite: ptr::null_mut(),
384            pclose: ptr::null_mut(),
385            perror: ptr::null_mut(),
386            popen: ptr::null_mut(),
387            puts: ptr::null_mut(),
388            rewind: ptr::null_mut(),
389            scanf: ptr::null_mut(),
390            setbuf: ptr::null_mut(),
391            setvbuf: ptr::null_mut(),
392            sscanf: ptr::null_mut(),
393            tempnam: ptr::null_mut(),
394            tmpfile: ptr::null_mut(),
395            tmpnam: ptr::null_mut(),
396            ungetc: ptr::null_mut(),
397            ai: ptr::null_mut(),
398            getenv: ptr::null_mut(),
399            breakfunc: ptr::null_mut(),
400            breakarg: ptr::null_mut(),
401            snprintf: ptr::null_mut(),
402            vsnprintf: ptr::null_mut(),
403            addrand: ptr::null_mut(),
404            addrandinit: Some(trampoline_addrandinit),
405        });
406
407        // Drive registrations into a thread-local sink so the C trampoline
408        // has somewhere to deposit them without capturing Rust state.
409        REGISTRY_SINK.with(|sink| {
410            let mut guard = sink.borrow_mut();
411            assert!(
412                guard.is_none(),
413                "nested ExternalLibrary::load is not supported"
414            );
415            *guard = Some(HashMap::new());
416        });
417
418        // SAFETY: funcadd is a valid C function from the loaded library; we
419        // pass it a correctly-shaped AmplExports.
420        unsafe { funcadd(ae.as_mut()) };
421
422        let funcs = REGISTRY_SINK
423            .with(|sink| sink.borrow_mut().take())
424            .unwrap_or_default();
425
426        Ok(ExternalLibrary {
427            _lib: Arc::new(lib),
428            _ae: ae,
429            funcs,
430        })
431    }
432
433    /// Names of all functions registered by this library.
434    pub fn function_names(&self) -> impl Iterator<Item = &str> {
435        self.funcs.keys().map(|s| s.as_str())
436    }
437
438    /// Look up a registered function by name.
439    pub fn get(&self, name: &str) -> Option<&RegisteredFunc> {
440        self.funcs.get(name)
441    }
442
443    /// Evaluate a registered function with the given positional arguments.
444    ///
445    /// Arguments are encoded per the AMPL `arglist` ABI: real args are stored
446    /// in `ra[]`, string args in `sa[]`, and `at[i]` maps argument position
447    /// `i` to either a real-slot index (`at[i] >= 0`) or a string-slot index
448    /// (`at[i] < 0`, decoded as `-(at[i]+1)`).
449    ///
450    /// If `want_derivs` is set, a length-`nr` derivative buffer is allocated
451    /// and returned on success. If `want_hes` is set, a length-`nr*(nr+1)/2`
452    /// Hessian buffer is also allocated and returned. The library is told to
453    /// fill both by the non-null `arglist.derivs` / `arglist.hes` pointers.
454    pub fn eval(
455        &self,
456        name: &str,
457        args: &[ExternalArg<'_>],
458        want_derivs: bool,
459        want_hes: bool,
460    ) -> Result<EvalResult, String> {
461        let rf = self
462            .funcs
463            .get(name)
464            .ok_or_else(|| format!("no such external function '{name}'"))?;
465
466        // Validate arity against the registered signature.
467        let n = args.len() as i32;
468        if rf.nargs >= 0 {
469            if rf.nargs != n {
470                return Err(format!(
471                    "external '{name}' expects {} args, got {}",
472                    rf.nargs, n
473                ));
474            }
475        } else {
476            // Negative: minimum -(nargs+1) args.
477            let min_args = -(rf.nargs + 1);
478            if n < min_args {
479                return Err(format!(
480                    "external '{name}' expects at least {min_args} args, got {n}"
481                ));
482            }
483        }
484
485        // Bucket args: build at[], ra[], sa[] in lockstep with their indices.
486        let mut at_vec: Vec<c_int> = Vec::with_capacity(args.len());
487        let mut ra_vec: Vec<f64> = Vec::new();
488        let mut sa_owned: Vec<CString> = Vec::new();
489        for a in args {
490            match a {
491                ExternalArg::Real(x) => {
492                    at_vec.push(ra_vec.len() as c_int);
493                    ra_vec.push(*x);
494                }
495                ExternalArg::Str(s) => {
496                    let cs = CString::new(*s)
497                        .map_err(|_| format!("external '{name}' string arg contains NUL"))?;
498                    at_vec.push(-(sa_owned.len() as c_int + 1));
499                    sa_owned.push(cs);
500                }
501            }
502        }
503        let nr = ra_vec.len() as c_int;
504        let sa_ptrs: Vec<*const c_char> = sa_owned.iter().map(|s| s.as_ptr()).collect();
505
506        // If the library declared FUNCADD_STRING_ARGS we let it see sa; if it
507        // did not, the library shouldn't be called with strings. Surface that.
508        let has_strings = !sa_owned.is_empty();
509        if has_strings && (rf.ty & FUNCADD_STRING_ARGS) == 0 {
510            return Err(format!(
511                "external '{name}' is not declared FUNCADD_STRING_ARGS but was \
512                 called with string arguments"
513            ));
514        }
515
516        // Optional output buffers.
517        let mut derivs_buf: Vec<f64> = if want_derivs {
518            vec![0.0; nr as usize]
519        } else {
520            Vec::new()
521        };
522        let hes_len = if want_hes {
523            (nr as usize) * ((nr as usize) + 1) / 2
524        } else {
525            0
526        };
527        let mut hes_buf: Vec<f64> = if want_hes {
528            vec![0.0; hes_len]
529        } else {
530            Vec::new()
531        };
532
533        // Space for a library-set error message.
534        let mut errmsg_buf: Vec<c_char> = vec![0; 1024];
535
536        // Build the arglist. Pointers into Rust-owned buffers are valid for
537        // the duration of the call since we hold those Vecs in this stack
538        // frame and the callee runs synchronously.
539        let mut al = Arglist {
540            n,
541            nr,
542            at: if at_vec.is_empty() {
543                ptr::null_mut()
544            } else {
545                at_vec.as_mut_ptr()
546            },
547            ra: if ra_vec.is_empty() {
548                ptr::null_mut()
549            } else {
550                ra_vec.as_mut_ptr()
551            },
552            sa: if sa_ptrs.is_empty() {
553                ptr::null_mut()
554            } else {
555                sa_ptrs.as_ptr() as *mut *const c_char
556            },
557            derivs: if want_derivs {
558                derivs_buf.as_mut_ptr()
559            } else {
560                ptr::null_mut()
561            },
562            hes: if want_hes {
563                hes_buf.as_mut_ptr()
564            } else {
565                ptr::null_mut()
566            },
567            dig: ptr::null_mut(),
568            funcinfo: rf.funcinfo,
569            // Some libraries read arglist.ae (e.g. to call fprintf); point at
570            // the same AmplExports we handed to funcadd_ASL.
571            ae: self._ae_ptr(),
572            f: ptr::null_mut(),
573            tva: ptr::null_mut(),
574            errmsg: errmsg_buf.as_mut_ptr(),
575            tmi: ptr::null_mut(),
576            private: ptr::null_mut(),
577            nin: 0,
578            nout: 0,
579            nsin: 0,
580            nsout: 0,
581        };
582
583        // SAFETY: rfunc is a valid extern "C" function pointer provided by
584        // the loaded library; arglist layout matches funcadd.h exactly.
585        // The AMPL lock serialises concurrent entry into the library.
586        let _guard = ampl_lock().lock().unwrap_or_else(|e| e.into_inner());
587        let value = unsafe { (rf.rfunc)(&mut al as *mut Arglist) };
588        drop(_guard);
589
590        // If the library wrote into errmsg, surface that. AMPL convention:
591        // if errmsg[0] != 0 after the call, treat as error.
592        if errmsg_buf[0] != 0 {
593            // SAFETY: errmsg_buf is a NUL-terminated C buffer (we allocated
594            // and zeroed it); the library only writes a C string there.
595            let msg = unsafe { CStr::from_ptr(errmsg_buf.as_ptr()) }
596                .to_string_lossy()
597                .into_owned();
598            return Err(format!("external '{name}' reported: {msg}"));
599        }
600
601        Ok(EvalResult {
602            value,
603            derivs: if want_derivs { Some(derivs_buf) } else { None },
604            hessian: if want_hes { Some(hes_buf) } else { None },
605        })
606    }
607
608    // Raw mutable pointer to the owned AmplExports. Used when building an
609    // arglist so the library can call back through the same table it was
610    // registered with. The Box is pinned for the lifetime of self.
611    fn _ae_ptr(&self) -> *mut AmplExports {
612        // Cast away the const; we never mutate the AmplExports ourselves.
613        (&*self._ae as *const AmplExports) as *mut AmplExports
614    }
615}
616
617/// One positional argument to an external function.
618#[derive(Debug, Clone, Copy)]
619pub enum ExternalArg<'a> {
620    Real(f64),
621    Str(&'a str),
622}
623
624/// Return value from [`ExternalLibrary::eval`].
625#[derive(Debug, Clone)]
626pub struct EvalResult {
627    /// Function value.
628    pub value: f64,
629    /// `df/dx_i` for each real argument, in `ra[]` order, if `want_derivs`.
630    pub derivs: Option<Vec<f64>>,
631    /// Packed upper-triangular Hessian in AMPL's convention,
632    /// `hes[i + j*(j+1)/2]` for `0 <= i <= j < nr`, if `want_hes`.
633    pub hessian: Option<Vec<f64>>,
634}
635
636// ---------------------------------------------------------------------------
637// Registration trampoline.
638//
639// `funcadd_ASL` can call Addfunc multiple times (once per registered name).
640// Rust closures can't be converted to `extern "C"` function pointers, so we
641// route each call through a free function that deposits into a thread-local
642// sink populated by `ExternalLibrary::load`.
643// ---------------------------------------------------------------------------
644
645thread_local! {
646    static REGISTRY_SINK: std::cell::RefCell<Option<HashMap<String, RegisteredFunc>>> =
647        std::cell::RefCell::new(None);
648}
649
650/// C-callable trampoline that receives Addfunc calls from the shared library.
651unsafe extern "C" fn trampoline_addfunc(
652    name: *const c_char,
653    f: Rfunc,
654    ty: c_int,
655    nargs: c_int,
656    funcinfo: *mut c_void,
657    _ae: *mut AmplExports,
658) {
659    if name.is_null() {
660        return;
661    }
662    // SAFETY: AMPL guarantees name is a NUL-terminated C string.
663    let cname = unsafe { CStr::from_ptr(name) };
664    let name_str = match cname.to_str() {
665        Ok(s) => s.to_owned(),
666        Err(_) => return, // non-UTF8 name — skip; real libs use ASCII.
667    };
668    REGISTRY_SINK.with(|sink| {
669        if let Some(map) = sink.borrow_mut().as_mut() {
670            map.insert(
671                name_str.clone(),
672                RegisteredFunc {
673                    name: name_str,
674                    rfunc: f,
675                    ty: ty as i32,
676                    nargs: nargs as i32,
677                    funcinfo,
678                },
679            );
680        }
681    });
682}
683
684/// Stub — some libraries ask us to register an AtReset callback. Pyomo logs a
685/// warning and does nothing. We do the same.
686unsafe extern "C" fn trampoline_atreset(_ae: *mut AmplExports, _f: *mut c_void, _v: *mut c_void) {
687    tracing::debug!("external library registered an AtReset callback; ignoring");
688}
689
690/// Stub — invoked by libraries that use random-valued externals. We just
691/// seed with 1 (matches Pyomo's default; no randomness in KKT paths).
692unsafe extern "C" fn trampoline_addrandinit(
693    _ae: *mut AmplExports,
694    setter: RandSeedSetter,
695    v: *mut c_void,
696) {
697    unsafe { setter(v, 1) };
698}
699
700#[cfg(test)]
701mod tests {
702    use super::*;
703
704    fn idaes_dylib() -> Option<std::path::PathBuf> {
705        let home = std::env::var_os("HOME")?;
706        let p = std::path::PathBuf::from(home).join(".idaes/bin/general_helmholtz_external.dylib");
707        if p.exists() {
708            Some(p)
709        } else {
710            None
711        }
712    }
713
714    fn idaes_params_dir() -> Option<String> {
715        let home = std::env::var_os("HOME")?;
716        let p = std::path::PathBuf::from(home).join(
717            "Dropbox/uv/.venv/lib/python3.12/site-packages/idaes/\
718             models/properties/general_helmholtz/components/parameters/",
719        );
720        if p.exists() {
721            p.to_str().map(|s| s.to_owned())
722        } else {
723            None
724        }
725    }
726
727    /// Opening the IDAES Helmholtz dylib (when present locally) should
728    /// surface the three functions used by the issue #15 fixture.
729    #[test]
730    fn load_idaes_helmholtz_dylib_registers_known_functions() {
731        let Some(path) = idaes_dylib() else {
732            eprintln!("skipping: IDAES dylib not present");
733            return;
734        };
735
736        let lib = ExternalLibrary::load(&path).expect("load should succeed");
737        let names: Vec<String> = lib.function_names().map(|s| s.to_owned()).collect();
738
739        for required in &["vf_hp", "h_liq_hp", "h_vap_hp"] {
740            assert!(
741                names.iter().any(|n| n == required),
742                "expected {required} in registered names: {names:?}"
743            );
744        }
745    }
746
747    /// Evaluate vf_hp at the NL fixture's initial guess. We don't assert the
748    /// exact numeric value (that's an IDAES invariant, not a ripopt one), but
749    /// the return value must be finite and the call must not set errmsg.
750    #[test]
751    fn eval_vf_hp_at_fixture_initial_point() {
752        let Some(path) = idaes_dylib() else {
753            eprintln!("skipping: IDAES dylib not present");
754            return;
755        };
756        let Some(params_dir) = idaes_params_dir() else {
757            eprintln!("skipping: IDAES parameters directory not present");
758            return;
759        };
760
761        let lib = ExternalLibrary::load(&path).expect("load");
762        // Fixture initial guess: h = 1878.71 kJ/kg-scaled, p = 101.325 kPa
763        // (the scaled values actually passed through the v3/v4 slots are
764        // 1878.71 * 0.0555... and 101325 * 0.001 respectively; using raw
765        // values here, the function should still return a finite number).
766        let args = [
767            ExternalArg::Str("h2o"),
768            ExternalArg::Real(1878.71 * 0.055508472036052976),
769            ExternalArg::Real(101325.0 * 0.001),
770            ExternalArg::Str(&params_dir),
771        ];
772        let res = lib.eval("vf_hp", &args, false, false).expect("eval");
773        assert!(
774            res.value.is_finite(),
775            "vf_hp returned non-finite value {}",
776            res.value
777        );
778    }
779
780    /// Same call path, but asking for first derivatives. derivs must be a
781    /// length-2 buffer (nr=2) of finite values.
782    #[test]
783    fn eval_vf_hp_with_derivatives() {
784        let Some(path) = idaes_dylib() else {
785            eprintln!("skipping: IDAES dylib not present");
786            return;
787        };
788        let Some(params_dir) = idaes_params_dir() else {
789            eprintln!("skipping: IDAES parameters directory not present");
790            return;
791        };
792
793        let lib = ExternalLibrary::load(&path).expect("load");
794        let args = [
795            ExternalArg::Str("h2o"),
796            ExternalArg::Real(1878.71 * 0.055508472036052976),
797            ExternalArg::Real(101325.0 * 0.001),
798            ExternalArg::Str(&params_dir),
799        ];
800        let res = lib.eval("vf_hp", &args, true, false).expect("eval");
801        let derivs = res.derivs.expect("derivs requested");
802        assert_eq!(derivs.len(), 2, "nr=2 reals -> 2 derivatives");
803        for (i, d) in derivs.iter().enumerate() {
804            assert!(d.is_finite(), "derivs[{i}] = {d} not finite");
805        }
806    }
807
808    /// Also request the packed Hessian. For nr=2 reals, that's 3 entries
809    /// (H00, H01, H11) in AMPL's packed upper-triangular layout.
810    #[test]
811    fn eval_vf_hp_with_hessian() {
812        let Some(path) = idaes_dylib() else {
813            eprintln!("skipping: IDAES dylib not present");
814            return;
815        };
816        let Some(params_dir) = idaes_params_dir() else {
817            eprintln!("skipping: IDAES parameters directory not present");
818            return;
819        };
820
821        let lib = ExternalLibrary::load(&path).expect("load");
822        let args = [
823            ExternalArg::Str("h2o"),
824            ExternalArg::Real(1878.71 * 0.055508472036052976),
825            ExternalArg::Real(101325.0 * 0.001),
826            ExternalArg::Str(&params_dir),
827        ];
828        let res = lib.eval("vf_hp", &args, true, true).expect("eval");
829        let hes = res.hessian.expect("hessian requested");
830        assert_eq!(hes.len(), 3, "nr=2 -> packed Hessian of length 3");
831        for (i, h) in hes.iter().enumerate() {
832            assert!(h.is_finite(), "hes[{i}] = {h} not finite");
833        }
834    }
835}
pounce_nl/nl_external.rs

pounce_nl/
nl_external.rs
