Repository: MiSawa/xq (rust)

Issue:
Pass the input iterator to vm
Current design that passes input value one by one and return an iterator per input value is simple, but it's impossible to handle `inputs`. We need to either pass the input iterator itself to the vm so vm can actively eat values from input.


Pull request description:
Use input iterator
Resolves #109 

Patch:
diff --git a/src/bin/main.rs b/src/bin/main.rs
index e4e5397..8397fea 100644
--- a/src/bin/main.rs
+++ b/src/bin/main.rs
@@ -5,7 +5,7 @@ use std::{
     io::{stdin, stdout, Write},
     path::PathBuf,
 };
-use xq::{module_loader::PreludeLoader, run_query, Value};
+use xq::{module_loader::PreludeLoader, run_query, InputError, Value};
 
 #[derive(Parser, Debug)]
 #[clap(author, about, version)]
@@ -126,10 +126,10 @@ fn init_log(verbosity: &Verbosity) -> Result<()> {
     .with_context(|| "Unable to initialize logger")
 }
 
-fn main() -> Result<()> {
-    let args: MainArgs = MainArgs::parse();
-    init_log(&args.verbosity)?;
-    log::debug!("Parsed argument: {:?}", args);
+fn run_with_input(
+    args: MainArgs,
+    input: impl Iterator<Item = Result<Value, InputError>>,
+) -> Result<()> {
     let query = if let Some(path) = args.query_file {
         log::trace!("Read query from file {:?}", path);
         std::fs::read_to_string(path)?
@@ -140,69 +140,73 @@ fn main() -> Result<()> {
         );
         args.query
     };
+    let module_loader = PreludeLoader();
+
+    let result_iterator = run_query(&query, input, &module_loader)
+        .map_err(|e| anyhow!("{:?}", e))
+        .with_context(|| "compile query")?;
+
     let output_format = args.output_format.get();
-    let output = |value| -> Result<()> {
-        match value {
-            Ok(value) => match output_format {
-                OutputFormat::Json => match value {
-                    Value::String(s) if args.output_format.raw_output => stdout()
-                        .write_all(s.as_bytes())
-                        .with_context(|| "write to output"),
-                    _ => if args.output_format.compact_output {
-                        serde_json::ser::to_writer::<_, Value>(stdout(), &value)
-                    } else {
-                        serde_json::ser::to_writer_pretty::<_, Value>(stdout(), &value)
+    match output_format {
+        OutputFormat::Json => {
+            for value in result_iterator {
+                match value {
+                    Ok(Value::String(s)) if args.output_format.raw_output => {
+                        stdout().write_all(s.as_bytes())?;
+                        println!();
+                    }
+                    Ok(value) => {
+                        if args.output_format.compact_output {
+                            serde_json::ser::to_writer::<_, Value>(stdout().lock(), &value)?;
+                            println!();
+                        } else {
+                            serde_json::ser::to_writer_pretty::<_, Value>(stdout().lock(), &value)?;
+                            println!();
+                        }
                     }
-                    .with_context(|| "write to output"),
-                },
-                OutputFormat::Yaml => serde_yaml::to_writer::<_, Value>(stdout(), &value)
-                    .with_context(|| "Write to output"),
+                    Err(e) => eprintln!("Error: {e:?}"),
+                }
             }
-            .and_then(|()| writeln!(stdout()).with_context(|| "Write ln"))?,
-            Err(e) => {
-                eprintln!("Error: {:?}", e);
+        }
+        OutputFormat::Yaml => {
+            for value in result_iterator {
+                match value {
+                    Ok(value) => serde_yaml::to_writer::<_, Value>(stdout().lock(), &value)
+                        .with_context(|| "Write to output")?,
+                    Err(e) => eprintln!("Error: {e:?}"),
+                }
             }
         }
-        Ok(())
-    };
+    }
+    Ok(())
+}
+
+fn main() -> Result<()> {
+    let args: MainArgs = MainArgs::parse();
+    init_log(&args.verbosity)?;
+    log::debug!("Parsed argument: {:?}", args);
 
-    let module_loader = PreludeLoader();
     match args.input_format.get() {
         InputFormat::Null => {
-            let results = run_query(&query, vec![Value::Null].into_iter(), &module_loader)
-                .map_err(|e| anyhow!("{:?}", e))
-                .with_context(|| "compile query")?;
-            for value in results {
-                output(value)?;
-            }
+            run_with_input(args, [Ok(Value::Null)].into_iter())?;
         }
         InputFormat::Json => {
             let stdin = stdin();
             let locked = stdin.lock();
-            let input: Vec<_> = serde_json::de::Deserializer::from_reader(locked)
+            let input = serde_json::de::Deserializer::from_reader(locked)
                 .into_iter::<Value>()
-                .collect::<Result<_, serde_json::Error>>()?;
-            let results = run_query(&query, input.into_iter(), &module_loader)
-                .map_err(|e| anyhow!("{:?}", e))
-                .with_context(|| "compile query")?;
-            for value in results {
-                output(value)?;
-            }
+                .map(|r| r.map_err(InputError::new));
+            run_with_input(args, input)?;
         }
         InputFormat::Yaml => {
             use serde::Deserialize;
             let stdin = stdin();
             let locked = stdin.lock();
-            let input: Vec<_> = serde_yaml::Deserializer::from_reader(locked)
+            let input = serde_yaml::Deserializer::from_reader(locked)
                 .into_iter()
                 .map(Value::deserialize)
-                .collect::<Result<_, serde_yaml::Error>>()?;
-            let results = run_query(&query, input.into_iter(), &module_loader)
-                .map_err(|e| anyhow!("{:?}", e))
-                .with_context(|| "compile query")?;
-            for value in results {
-                output(value)?;
-            }
+                .map(|r| r.map_err(InputError::new));
+            run_with_input(args, input)?;
         }
     }
     Ok(())
diff --git a/src/lib.rs b/src/lib.rs
index 5760f23..cb455b4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -20,6 +20,9 @@ use crate::{
     vm::{machine::Machine, QueryExecutionError},
 };
 use thiserror::Error;
+use vm::machine::ResultIterator;
+
+pub type InputError = vm::error::InputError;
 
 #[derive(Debug, Error)]
 pub enum XQError {
@@ -35,9 +38,9 @@ pub fn run_query<I, M>(
     query: &str,
     input: I,
     module_loader: &M,
-) -> Result<OutputIterator<I>, XQError>
+) -> Result<ResultIterator<I>, XQError>
 where
-    I: Iterator<Item = Value>,
+    I: Iterator<Item = Result<Value, InputError>>,
     M: ModuleLoader,
 {
     // let now = std::time::Instant::now();
@@ -51,49 +54,6 @@ where
     log::info!("Compiled program = {:?}", program);
     // eprintln!("Compile: {:?}", now.elapsed());
 
-    let vm = Machine::new(program);
-    Ok(OutputIterator::new(vm, input))
-}
-
-pub struct OutputIterator<I: Iterator<Item = Value>> {
-    finished: bool,
-    machine: Machine,
-    results: Option<vm::machine::ResultIterator>,
-    input: I,
-}
-
-impl<I: Iterator<Item = Value>> OutputIterator<I> {
-    fn new(machine: Machine, input: I) -> Self {
-        Self {
-            finished: false,
-            machine,
-            results: None,
-            input,
-        }
-    }
-}
-
-impl<I: Iterator<Item = Value>> Iterator for OutputIterator<I> {
-    type Item = Result<Value, QueryExecutionError>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.finished {
-            return None;
-        }
-        loop {
-            if let Some(it) = &mut self.results {
-                if let Some(value) = it.next() {
-                    return Some(value);
-                } else {
-                    self.results = None;
-                }
-            }
-            if let Some(value) = self.input.next() {
-                self.results = Some(self.machine.run(value))
-            } else {
-                self.finished = true;
-                return None;
-            }
-        }
-    }
+    let mut vm = Machine::new(program);
+    Ok(vm.start(input))
 }
diff --git a/src/vm/error.rs b/src/vm/error.rs
index 159f7da..a519d20 100644
--- a/src/vm/error.rs
+++ b/src/vm/error.rs
@@ -5,7 +5,17 @@ use thiserror::Error;
 
 pub type Result<T, E = QueryExecutionError> = std::result::Result<T, E>;
 
-#[derive(Debug, Clone, Eq, PartialEq, Error)]
+#[derive(Debug, Error)]
+#[error(transparent)]
+pub struct InputError(#[from] Box<dyn std::error::Error>);
+impl InputError {
+    pub fn new<E: 'static + std::error::Error>(e: E) -> Self {
+        Self(Box::new(e))
+    }
+}
+
+#[derive(Debug, Error)]
+#[non_exhaustive]
 pub enum QueryExecutionError {
     #[error("Object was indexed by non-string value `{0:?}`")]
     ObjectIndexByNonString(Value),
@@ -73,6 +83,8 @@ pub enum QueryExecutionError {
     IndeterminateOffset(#[from] time::error::IndeterminateOffset),
     #[error("Unable to determine local time zone")]
     TimeZoneLookupFailure(RcString),
+    #[error("Input source gave an error")]
+    InputError(#[from] InputError),
     #[error("{0:?}")]
     UserDefinedError(Value),
 }
diff --git a/src/vm/machine.rs b/src/vm/machine.rs
index 4e14501..e2ca4c1 100644
--- a/src/vm/machine.rs
+++ b/src/vm/machine.rs
@@ -13,11 +13,12 @@ use crate::{
         error::QueryExecutionError,
         Address, ByteCode, Program, Result, ScopeId, ScopedSlot, Value,
     },
-    Array,
+    Array, InputError,
 };
 use itertools::Itertools;
 use std::{
     cell::{RefCell, RefMut},
+    iter::Fuse,
     rc::Rc,
 };
 use thiserror::Error;
@@ -158,13 +159,14 @@ enum OnFork {
     TryAlternative,
     CatchLabel(LabelId),
     Iterate,
+    IterateInput,
 }
 
 impl Environment {
     fn new(state: <State as Undo>::UndoToken) -> Self {
         Self {
             next_label_id: 0,
-            forks: vec![(state, OnFork::Nop)],
+            forks: vec![(state, OnFork::IterateInput)],
         }
     }
 
@@ -429,33 +431,47 @@ impl Machine {
         }
     }
 
-    pub fn run(&mut self, value: Value) -> ResultIterator {
+    pub fn start<I: Iterator<Item = Result<Value, InputError>>>(
+        &mut self,
+        input: I,
+    ) -> ResultIterator<I> {
         let mut state = State::new(self.program.entry_point);
-        state.push(value);
         let env = Environment::new(state.save());
         ResultIterator {
             program: self.program.clone(),
             env,
             state,
+            input: input.fuse(),
         }
     }
 }
 
-pub struct ResultIterator {
+pub struct ResultIterator<I: Iterator<Item = Result<Value, InputError>>> {
     program: Rc<Program>,
     env: Environment,
     state: State,
+    input: Fuse<I>,
 }
 
-impl Iterator for ResultIterator {
+impl<I: Iterator<Item = Result<Value, InputError>>> Iterator for ResultIterator<I> {
     type Item = Result<Value>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        run_code(&self.program, &mut self.state, &mut self.env)
+        run_code(
+            &self.program,
+            &mut self.state,
+            &mut self.env,
+            &mut self.input,
+        )
     }
 }
 
-fn run_code(program: &Program, state: &mut State, env: &mut Environment) -> Option<Result<Value>> {
+fn run_code(
+    program: &Program,
+    state: &mut State,
+    env: &mut Environment,
+    input: &mut impl Iterator<Item = Result<Value, InputError>>,
+) -> Option<Result<Value>> {
     let mut err: Option<QueryExecutionError> = None;
     log::trace!("Start from environment {:?}", env);
     'backtrack: loop {
@@ -539,6 +555,26 @@ fn run_code(program: &Program, state: &mut State, env: &mut Environment) -> Opti
                         }
                     }
                 }
+                (OnFork::IterateInput, None) => {
+                    match input.next() {
+                        None => return None,
+                        Some(Err(e)) => {
+                            // Don't push fork again so we don't infinitely see input errors
+                            // even when the input stream gave us infinite errors.
+                            return Some(Err(e.into()));
+                        }
+                        Some(Ok(v)) => {
+                            state.undo(token);
+                            env.push_fork(state, OnFork::IterateInput, state.pc);
+                            state.push(v);
+                            break 'select_fork state.save();
+                        }
+                    }
+                }
+                (OnFork::IterateInput, Some(_e)) => {
+                    env.push_fork(state, OnFork::IterateInput, state.pc);
+                    return Some(Err(err.take().unwrap()));
+                }
                 (_, Some(_)) => continue 'select_fork,
             }
             break 'select_fork token;


Test patch:
diff --git a/tests/common/mod.rs b/tests/common/mod.rs
index 1f39ce9..570c7d0 100644
--- a/tests/common/mod.rs
+++ b/tests/common/mod.rs
@@ -1,5 +1,5 @@
 use std::error::Error;
-use xq::{module_loader::PreludeLoader, run_query, Value};
+use xq::{module_loader::PreludeLoader, run_query, InputError, Value};
 
 #[macro_export]
 macro_rules! test {
@@ -12,9 +12,9 @@ macro_rules! test {
 }
 
 pub(crate) fn run_test(query: &str, input: &str, output: &str) -> Result<(), Box<dyn Error>> {
-    let input: Vec<_> = serde_json::de::Deserializer::from_str(input)
+    let input = serde_json::de::Deserializer::from_str(input)
         .into_iter::<Value>()
-        .collect::<Result<_, serde_json::Error>>()?;
+        .map(|r| r.map_err(InputError::new));
     let expected: Vec<_> = serde_json::de::Deserializer::from_str(output)
         .into_iter::<Value>()
         .collect::<Result<_, serde_json::Error>>()?;