1use crate::core::{self, QueryResult};
13use crate::machine::{Machine, RegistryEntry, SrcLoc};
14use crate::wire::{EncoderDesc, Envelope, WireError};
15use plg_shared::StringInterner;
16use std::ffi::CStr;
17use std::io::{self, Read, Write};
18use std::os::raw::c_char;
19
20#[unsafe(no_mangle)]
28pub unsafe extern "C" fn plg_rt_init(
29 atom_strs: *const *const c_char,
30 atom_count: u32,
31 registry: *const RegistryEntry,
32 registry_len: u32,
33 srcmap: *const SrcLoc,
34 srcmap_len: u32,
35 files: *const *const c_char,
36 files_len: u32,
37 caps: *const *const crate::wire::EncoderDesc,
38 caps_len: u32,
39) -> *mut Machine {
40 let mut atoms = StringInterner::new();
41 for i in 0..atom_count as usize {
42 let s = unsafe { CStr::from_ptr(*atom_strs.add(i)) };
43 let expected = i as u32;
44 let id = atoms.intern(&s.to_string_lossy());
45 debug_assert_eq!(id, expected, "atom table out of sync with interner");
46 }
47 let registry: Vec<RegistryEntry> =
48 unsafe { std::slice::from_raw_parts(registry, registry_len as usize) }.to_vec();
49 debug_assert!(
50 registry.is_sorted_by_key(|e| (e.functor, e.arity)),
51 "registry must be sorted for binary search"
52 );
53 let srcmap: Vec<SrcLoc> = if srcmap_len == 0 {
54 Vec::new()
55 } else {
56 unsafe { std::slice::from_raw_parts(srcmap, srcmap_len as usize) }.to_vec()
57 };
58 let files: Vec<String> = (0..files_len as usize)
59 .map(|i| {
60 unsafe { CStr::from_ptr(*files.add(i)) }
61 .to_string_lossy()
62 .into_owned()
63 })
64 .collect();
65 let mut m = Machine::new(atoms, registry);
66 m.set_provenance(srcmap, files);
67 m.capabilities = (0..caps_len as usize)
68 .map(|i| unsafe { *caps.add(i) })
69 .collect();
70 Box::into_raw(m)
71}
72
73struct Args {
74 query: Option<String>,
75 limit: Option<usize>,
76 format: String,
77 input_format: String,
78 atoms: bool,
79}
80
81fn parse_args(argv: Vec<String>) -> Result<Args, String> {
82 let mut query = None;
83 let mut limit = None;
84 let mut format = "text".to_string(); let mut input_format = "text".to_string(); let mut atoms = false; let mut it = argv.into_iter().peekable();
88 while let Some(arg) = it.next() {
89 let (flag, inline_value) = match arg.split_once('=') {
90 Some((f, v)) => (f.to_string(), Some(v.to_string())),
91 None => (arg, None),
92 };
93 let value = |it: &mut std::iter::Peekable<std::vec::IntoIter<String>>| {
94 inline_value
95 .clone()
96 .or_else(|| it.next())
97 .ok_or(format!("missing value for {flag}"))
98 };
99 match flag.as_str() {
100 "-q" | "--query" => query = Some(value(&mut it)?),
101 "-l" | "--limit" => {
102 limit = Some(
103 value(&mut it)?
104 .parse::<usize>()
105 .map_err(|_| "invalid --limit value".to_string())?,
106 )
107 }
108 "-f" | "--format" => format = value(&mut it)?,
109 "--input-format" => input_format = value(&mut it)?,
110 "--atoms" => atoms = true,
111 "-h" | "--help" => {
112 return Err(
113 "usage: --query <goal> [--limit N] [--format text|bson] [--input-format text|bson] [--atoms]"
114 .to_string(),
115 );
116 }
117 other => return Err(format!("unexpected argument: {other}")),
118 }
119 }
120 Ok(Args {
121 query,
122 limit,
123 format,
124 input_format,
125 atoms,
126 })
127}
128
129fn output_solutions(enc: &EncoderDesc, m: &Machine, exhausted: bool, atoms: bool) {
135 let mut out = io::stdout().lock();
136 let mut env = Envelope::from_machine(m, exhausted);
137 let atom_names: Vec<String> = if atoms && !(enc.can_stream)() {
141 (0..m.atoms.len())
142 .map(|i| {
143 m.atoms
144 .try_resolve(i as u32)
145 .unwrap_or_default()
146 .to_string()
147 })
148 .collect()
149 } else {
150 Vec::new()
151 };
152 if !atom_names.is_empty() {
153 env.atoms = Some(&atom_names);
154 }
155 let _ = (enc.write_envelope)(&mut out, m, &env);
156 if !(enc.can_stream)() {
157 let _ = out.flush();
158 }
159}
160
161fn output_result(enc: Option<&EncoderDesc>, err: WireError) {
165 let message = match &err {
166 WireError::Parse(m) | WireError::Runtime(m) => m.as_str(),
167 };
168 match enc {
169 Some(e) => {
170 let mut out = io::stdout().lock();
171 let _ = (e.write_error)(&mut out, &err);
172 if !(e.can_stream)() {
173 let _ = out.flush();
174 }
175 }
176 None => eprintln!("Error: {message}"),
177 }
178}
179
180#[unsafe(no_mangle)]
186pub unsafe extern "C" fn plg_rt_main(
187 m: *mut Machine,
188 argc: i32,
189 argv: *const *const c_char,
190) -> i32 {
191 let m = unsafe { &mut *m };
192 let raw_args: Vec<String> = (1..argc as usize)
193 .map(|i| {
194 unsafe { CStr::from_ptr(*argv.add(i)) }
195 .to_string_lossy()
196 .into_owned()
197 })
198 .collect();
199
200 let args = match parse_args(raw_args) {
201 Ok(a) => a,
202 Err(e) => {
203 eprintln!("{e}");
204 return 2;
205 }
206 };
207 let enc: Option<&'static EncoderDesc> = match unsafe {
212 EncoderDesc::find(m.capabilities.as_ptr(), m.capabilities.len(), &args.format)
213 } {
214 Some(e) => Some(e),
215 None => {
216 eprintln!("Unknown or undeclared format: {}", args.format);
217 return 2;
218 }
219 };
220 if args.atoms && args.query.is_none() && args.input_format == "text" {
224 let mut out = io::stdout().lock();
225 let e = enc.unwrap();
226 if (e.can_stream)() {
227 let _ = crate::wire::write_atom_map_text(&mut out, m);
228 } else {
229 let _ = crate::wire::write_atom_map_bson(&mut out, m);
230 let _ = out.flush();
231 }
232 return 0;
233 }
234 let (query, argv_limit) = match args.input_format.as_str() {
240 "text" => {
241 let q = match args.query {
242 Some(q) => q,
243 None => {
244 eprintln!("missing required argument: --query <goal>");
245 return 2;
246 }
247 };
248 (q, args.limit)
249 }
250 "bson" => {
251 if unsafe {
252 EncoderDesc::find(m.capabilities.as_ptr(), m.capabilities.len(), "bson").is_none()
253 } {
254 eprintln!("Unknown or undeclared input format: bson");
255 return 2;
256 }
257 let mut stdin_buf = Vec::new();
258 if let Err(e) = std::io::stdin().read_to_end(&mut stdin_buf) {
259 output_result(
260 enc,
261 WireError::Parse(format!("bson request read error: {e}")),
262 );
263 return 2;
264 }
265 match crate::wire::parse_bson_request(&stdin_buf) {
266 Ok(req) => (req.query, args.limit.or(req.limit)),
267 Err(e) => {
268 output_result(
272 enc,
273 WireError::Parse(format!("bson request parse error: {e}")),
274 );
275 return 2;
276 }
277 }
278 }
279 other => {
280 eprintln!("Unknown --input-format: {other} (expected text|bson)");
281 return 2;
282 }
283 };
284 if let Some(e) = enc
288 && !(e.can_stream)()
289 {
290 m.output = crate::machine::OutputSink::Capture(String::new());
291 }
292 m.solution_limit = argv_limit;
293 if let Ok(s) = std::env::var("PLG_MAX_STEPS")
294 && let Ok(n) = s.parse::<u64>()
295 {
296 m.step_limit = n;
297 }
298 if let Ok(s) = std::env::var("PLG_METACALL_DEPTH")
299 && let Ok(n) = s.parse::<usize>()
300 {
301 m.metacall_depth_limit = n;
302 }
303
304 match core::run_query(m, &query) {
305 QueryResult::ParseError(msg) => {
306 output_result(enc, WireError::Parse(msg));
307 2
308 }
309 QueryResult::RuntimeError(msg) => {
310 output_result(enc, WireError::Runtime(msg));
311 3
312 }
313 QueryResult::Solutions => {
314 let count = m.solutions.len();
315 let exhausted = core::exhausted(m);
316 output_solutions(enc.unwrap(), m, exhausted, args.atoms);
317 if count > 0 { 1 } else { 0 }
318 }
319 }
320}
321
322#[cfg(test)]
323mod tests {
324 use super::*;
325
326 fn args(v: &[&str]) -> Result<Args, String> {
327 parse_args(v.iter().map(|s| s.to_string()).collect())
328 }
329
330 #[test]
331 fn parses_flags_with_space_and_equals() {
332 let a = args(&["--query", "p(X)", "--limit", "3", "--format", "bson"]).unwrap();
333 assert_eq!(a.query.as_deref(), Some("p(X)"));
334 assert_eq!(a.limit, Some(3));
335 assert_eq!(a.format, "bson");
336 assert_eq!(a.input_format, "text", "default input-format is text");
337
338 let a = args(&["--query=p(X)", "-l", "1"]).unwrap();
339 assert_eq!(a.query.as_deref(), Some("p(X)"));
340 assert_eq!(a.limit, Some(1));
341 assert_eq!(a.format, "text", "default format is text");
342 }
343
344 #[test]
345 fn parses_input_format_flag() {
346 let a = args(&["--query", "p(X)", "--input-format", "bson"]).unwrap();
347 assert_eq!(a.input_format, "bson");
348 assert!(args(&["--input-format", "bson"]).is_ok());
351 }
352
353 #[test]
354 fn missing_value_flags_are_errors() {
355 assert!(args(&["--query"]).is_err());
356 assert!(args(&["--bogus", "x"]).is_err());
357 assert!(args(&["--input-format"]).is_err());
358 }
359}