Repository: MiSawa/xq (rust)

Issue:
Rework on `-n`
https://github.com/MiSawa/xq/issues/116#issuecomment-1066038523
```bash
$ jq -nc 'input' <<< '1 2 3'
1
$ jq -nsc 'input' <<< '1 2 3'
[1,2,3]
```
I guess we need two streams, one for the things supplied to the program, and another for `input/0`, where both tied to a single shared iterator usually but does something different for `-n`?

Pull request description:
Fix how -n works
Fixes #120

Patch:
diff --git a/src/bin/cli/input.rs b/src/bin/cli/input.rs
new file mode 100644
index 0000000..6bf6309
--- /dev/null
+++ b/src/bin/cli/input.rs
@@ -0,0 +1,84 @@
+use std::iter::Once;
+
+use xq::{util::SharedIterator, Array, InputError, Value};
+
+type ResultValue = Result<Value, InputError>;
+
+pub(crate) trait Input {
+    type SingleInputIterator: Iterator<Item = ResultValue>;
+    type ContextsIterator: Iterator<Item = ResultValue>;
+    type InputIterator: Iterator<Item = ResultValue>;
+
+    fn into_input_iterator(self) -> Self::SingleInputIterator;
+    fn into_iterators(self) -> (Self::ContextsIterator, Self::InputIterator);
+
+    fn null_input(self) -> Null<Self>
+    where
+        Self: Sized,
+    {
+        Null(self)
+    }
+}
+
+pub(crate) struct Tied<I: Iterator<Item = ResultValue>>(I);
+impl<I: Iterator<Item = ResultValue>> Tied<I> {
+    pub(crate) fn new(inner: I) -> Self {
+        Self(inner)
+    }
+    pub(crate) fn slurp(self) -> impl Input {
+        let v = self
+            .0
+            .collect::<Result<Array, InputError>>()
+            .map(Into::into);
+        Slurp(v)
+    }
+}
+
+impl<I: Iterator<Item = ResultValue>> Input for Tied<I> {
+    type SingleInputIterator = I;
+    type ContextsIterator = SharedIterator<I>;
+    type InputIterator = SharedIterator<I>;
+
+    fn into_input_iterator(self) -> Self::SingleInputIterator {
+        self.0
+    }
+
+    fn into_iterators(self) -> (Self::ContextsIterator, Self::InputIterator) {
+        let shared = SharedIterator::from(self.0);
+        (shared.clone(), shared)
+    }
+}
+
+pub(crate) struct Null<I>(I);
+impl<I: Input> Input for Null<I> {
+    type SingleInputIterator = I::SingleInputIterator;
+    type ContextsIterator = Once<ResultValue>;
+    type InputIterator = I::SingleInputIterator;
+
+    fn into_input_iterator(self) -> Self::SingleInputIterator {
+        self.0.into_input_iterator()
+    }
+
+    fn into_iterators(self) -> (Self::ContextsIterator, Self::InputIterator) {
+        (
+            std::iter::once(Ok(Value::Null)),
+            self.0.into_input_iterator(),
+        )
+    }
+}
+
+struct Slurp(ResultValue);
+impl Input for Slurp {
+    type SingleInputIterator = Once<ResultValue>;
+    type ContextsIterator = SharedIterator<Once<ResultValue>>;
+    type InputIterator = SharedIterator<Once<ResultValue>>;
+
+    fn into_input_iterator(self) -> Self::SingleInputIterator {
+        std::iter::once(self.0)
+    }
+
+    fn into_iterators(self) -> (Self::ContextsIterator, Self::InputIterator) {
+        let shared = SharedIterator::from(std::iter::once(self.0));
+        (shared.clone(), shared)
+    }
+}
diff --git a/src/bin/cli/mod.rs b/src/bin/cli/mod.rs
new file mode 100644
index 0000000..489a2dd
--- /dev/null
+++ b/src/bin/cli/mod.rs
@@ -0,0 +1,1 @@
+pub(crate) mod input;
diff --git a/src/bin/main.rs b/src/bin/main.rs
index 70b4382..090e3e5 100644
--- a/src/bin/main.rs
+++ b/src/bin/main.rs
@@ -1,11 +1,16 @@
 use anyhow::{anyhow, Context, Result};
 use clap::{ArgEnum, Args, Parser, ValueHint};
 use clap_verbosity_flag::Verbosity;
+use cli::input::Input;
 use std::{
     io::{stdin, stdout, Write},
     path::PathBuf,
 };
-use xq::{module_loader::PreludeLoader, run_query, Array, InputError, Value};
+use xq::{module_loader::PreludeLoader, run_query, InputError, Value};
+
+use crate::cli::input::Tied;
+
+mod cli;
 
 #[derive(Parser, Debug)]
 #[clap(author, about, version)]
@@ -37,8 +42,7 @@ struct MainArgs {
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, ArgEnum)]
-enum InputFormat {
-    Null,
+enum SerializationFormat {
     Json,
     Yaml,
 }
@@ -46,11 +50,8 @@ enum InputFormat {
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Args)]
 struct InputFormatArg {
     /// Specify input format
-    #[clap(group = "input", long, arg_enum, default_value_t = InputFormat::Json)]
-    input_format: InputFormat,
-    /// Use null as an input value
-    #[clap(group = "input", short = 'n', long)]
-    null_input: bool,
+    #[clap(group = "input", long, arg_enum, default_value_t = SerializationFormat::Json)]
+    input_format: SerializationFormat,
     /// Read input as json values
     #[clap(group = "input", long)]
     json_input: bool,
@@ -58,36 +59,32 @@ struct InputFormatArg {
     #[clap(group = "input", long)]
     yaml_input: bool,
 
+    /// Single null is supplied to the program.
+    /// The original input can still be read via input/0 and inputs/0.
+    #[clap(short, long)]
+    null_input: bool,
     /// Read input values into an array
     #[clap(short, long)]
     slurp: bool,
 }
 
 impl InputFormatArg {
-    fn get(&self) -> InputFormat {
-        if self.null_input {
-            InputFormat::Null
-        } else if self.json_input {
-            InputFormat::Json
+    fn get(&self) -> SerializationFormat {
+        if self.json_input {
+            SerializationFormat::Json
         } else if self.yaml_input {
-            InputFormat::Yaml
+            SerializationFormat::Yaml
         } else {
             self.input_format
         }
     }
 }
 
-#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, ArgEnum)]
-enum OutputFormat {
-    Json,
-    Yaml,
-}
-
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Args)]
 struct OutputFormatArg {
     /// Specify output format
-    #[clap(group = "output", long, arg_enum, default_value_t = OutputFormat::Json)]
-    output_format: OutputFormat,
+    #[clap(group = "output", long, arg_enum, default_value_t = SerializationFormat::Json)]
+    output_format: SerializationFormat,
     /// Write output as json values
     #[clap(group = "output", long)]
     json_output: bool,
@@ -104,11 +101,11 @@ struct OutputFormatArg {
 }
 
 impl OutputFormatArg {
-    fn get(&self) -> OutputFormat {
+    fn get(&self) -> SerializationFormat {
         if self.json_output {
-            OutputFormat::Json
+            SerializationFormat::Json
         } else if self.yaml_output {
-            OutputFormat::Yaml
+            SerializationFormat::Yaml
         } else {
             self.output_format
         }
@@ -130,10 +127,7 @@ fn init_log(verbosity: &Verbosity) -> Result<()> {
     .with_context(|| "Unable to initialize logger")
 }
 
-fn run_with_input(
-    args: MainArgs,
-    input: impl Iterator<Item = Result<Value, InputError>>,
-) -> Result<()> {
+fn run_with_input(args: MainArgs, input: impl Input) -> Result<()> {
     let query = if let Some(path) = args.query_file {
         log::trace!("Read query from file {:?}", path);
         std::fs::read_to_string(path)?
@@ -146,13 +140,14 @@ fn run_with_input(
     };
     let module_loader = PreludeLoader();
 
-    let result_iterator = run_query(&query, input, &module_loader)
+    let (context, input) = input.into_iterators();
+    let result_iterator = run_query(&query, context, input, &module_loader)
         .map_err(|e| anyhow!("{:?}", e))
         .with_context(|| "compile query")?;
 
     let output_format = args.output_format.get();
     match output_format {
-        OutputFormat::Json => {
+        SerializationFormat::Json => {
             for value in result_iterator {
                 match value {
                     Ok(Value::String(s)) if args.output_format.raw_output => {
@@ -172,7 +167,7 @@ fn run_with_input(
                 }
             }
         }
-        OutputFormat::Yaml => {
+        SerializationFormat::Yaml => {
             for value in result_iterator {
                 match value {
                     Ok(value) => serde_yaml::to_writer::<_, Value>(stdout().lock(), &value)
@@ -185,29 +180,34 @@ fn run_with_input(
     Ok(())
 }
 
+fn run_with_raw_input<I: Iterator<Item = Result<Value, InputError>>>(
+    args: MainArgs,
+    i: I,
+) -> Result<()> {
+    let i = Tied::new(i);
+    match (args.input_format.slurp, args.input_format.null_input) {
+        (false, false) => run_with_input(args, i),
+        (false, true) => run_with_input(args, i.null_input()),
+        (true, false) => run_with_input(args, i.slurp()),
+        (true, true) => run_with_input(args, i.slurp().null_input()),
+    }
+}
+
 fn main() -> Result<()> {
     let args: MainArgs = MainArgs::parse();
     init_log(&args.verbosity)?;
     log::debug!("Parsed argument: {:?}", args);
 
     match args.input_format.get() {
-        InputFormat::Null => {
-            run_with_input(args, [Ok(Value::Null)].into_iter())?;
-        }
-        InputFormat::Json => {
+        SerializationFormat::Json => {
             let stdin = stdin();
             let locked = stdin.lock();
             let input = serde_json::de::Deserializer::from_reader(locked)
                 .into_iter::<Value>()
                 .map(|r| r.map_err(InputError::new));
-            if args.input_format.slurp {
-                let v: Array = input.collect::<Result<_, InputError>>()?;
-                run_with_input(args, [Ok(v.into())].into_iter())?;
-            } else {
-                run_with_input(args, input)?;
-            }
+            run_with_raw_input(args, input)
         }
-        InputFormat::Yaml => {
+        SerializationFormat::Yaml => {
             use serde::Deserialize;
             let stdin = stdin();
             let locked = stdin.lock();
@@ -215,13 +215,7 @@ fn main() -> Result<()> {
                 .into_iter()
                 .map(Value::deserialize)
                 .map(|r| r.map_err(InputError::new));
-            if args.input_format.slurp {
-                let v: Array = input.collect::<Result<_, InputError>>()?;
-                run_with_input(args, [Ok(v.into())].into_iter())?;
-            } else {
-                run_with_input(args, input)?;
-            }
+            run_with_raw_input(args, input)
         }
     }
-    Ok(())
 }
diff --git a/src/lib.rs b/src/lib.rs
index cb455b4..9416327 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -4,7 +4,7 @@ mod intrinsic;
 pub mod lang;
 pub mod module_loader;
 mod number;
-mod util;
+pub mod util;
 mod value;
 pub mod vm;
 
@@ -34,12 +34,14 @@ pub enum XQError {
     QueryExecutionError(#[from] QueryExecutionError),
 }
 
-pub fn run_query<I, M>(
+pub fn run_query<C, I, M>(
     query: &str,
+    context: C,
     input: I,
     module_loader: &M,
-) -> Result<ResultIterator<I>, XQError>
+) -> Result<ResultIterator<C, I>, XQError>
 where
+    C: Iterator<Item = Result<Value, InputError>>,
     I: Iterator<Item = Result<Value, InputError>>,
     M: ModuleLoader,
 {
@@ -55,5 +57,5 @@ where
     // eprintln!("Compile: {:?}", now.elapsed());
 
     let mut vm = Machine::new(program);
-    Ok(vm.start(input))
+    Ok(vm.start(context, input))
 }
diff --git a/src/util.rs b/src/util.rs
index 9557bbe..55389c7 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -1,8 +1,27 @@
-use std::rc::Rc;
+use std::{cell::RefCell, rc::Rc};
 
-pub fn make_owned<T: Clone>(v: Rc<T>) -> T {
+pub(crate) fn make_owned<T: Clone>(v: Rc<T>) -> T {
     match Rc::try_unwrap(v) {
         Ok(v) => v,
         Err(v) => (*v).clone(),
     }
 }
+
+pub struct SharedIterator<I>(Rc<RefCell<I>>);
+impl<I> From<I> for SharedIterator<I> {
+    fn from(it: I) -> Self {
+        Self(Rc::new(RefCell::new(it)))
+    }
+}
+impl<I> Clone for SharedIterator<I> {
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+impl<T, I: Iterator<Item = T>> Iterator for SharedIterator<I> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.borrow_mut().next()
+    }
+}
diff --git a/src/vm/machine.rs b/src/vm/machine.rs
index 5367b1f..8377edc 100644
--- a/src/vm/machine.rs
+++ b/src/vm/machine.rs
@@ -159,14 +159,14 @@ enum OnFork {
     TryAlternative,
     CatchLabel(LabelId),
     Iterate,
-    IterateInput,
+    IterateContext,
 }
 
 impl Environment {
     fn new(state: <State as Undo>::UndoToken) -> Self {
         Self {
             next_label_id: 0,
-            forks: vec![(state, OnFork::IterateInput)],
+            forks: vec![(state, OnFork::IterateContext)],
         }
     }
 
@@ -431,29 +431,42 @@ impl Machine {
         }
     }
 
-    pub fn start<I: Iterator<Item = Result<Value, InputError>>>(
+    pub fn start<
+        C: Iterator<Item = Result<Value, InputError>>,
+        I: Iterator<Item = Result<Value, InputError>>,
+    >(
         &mut self,
+        context: C,
         input: I,
-    ) -> ResultIterator<I> {
+    ) -> ResultIterator<C, I> {
         let mut state = State::new(self.program.entry_point);
         let env = Environment::new(state.save());
         ResultIterator {
             program: self.program.clone(),
             env,
             state,
+            context: context.fuse(),
             input: input.fuse(),
         }
     }
 }
 
-pub struct ResultIterator<I: Iterator<Item = Result<Value, InputError>>> {
+pub struct ResultIterator<
+    C: Iterator<Item = Result<Value, InputError>>,
+    I: Iterator<Item = Result<Value, InputError>>,
+> {
     program: Rc<Program>,
     env: Environment,
     state: State,
+    context: Fuse<C>,
     input: Fuse<I>,
 }
 
-impl<I: Iterator<Item = Result<Value, InputError>>> Iterator for ResultIterator<I> {
+impl<
+        C: Iterator<Item = Result<Value, InputError>>,
+        I: Iterator<Item = Result<Value, InputError>>,
+    > Iterator for ResultIterator<C, I>
+{
     type Item = Result<Value>;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -461,6 +474,7 @@ impl<I: Iterator<Item = Result<Value, InputError>>> Iterator for ResultIterator<
             &self.program,
             &mut self.state,
             &mut self.env,
+            &mut self.context,
             &mut self.input,
         )
     }
@@ -470,6 +484,7 @@ fn run_code(
     program: &Program,
     state: &mut State,
     env: &mut Environment,
+    context: &mut impl Iterator<Item = Result<Value, InputError>>,
     input: &mut impl Iterator<Item = Result<Value, InputError>>,
 ) -> Option<Result<Value>> {
     let mut err: Option<QueryExecutionError> = None;
@@ -555,24 +570,24 @@ fn run_code(
                         }
                     }
                 }
-                (OnFork::IterateInput, None) => {
-                    match input.next() {
+                (OnFork::IterateContext, None) => {
+                    match context.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.
+                            // even when the context stream gave us infinite errors.
                             return Some(Err(e.into()));
                         }
                         Some(Ok(v)) => {
                             state.undo(token);
-                            env.push_fork(state, OnFork::IterateInput, state.pc);
+                            env.push_fork(state, OnFork::IterateContext, state.pc);
                             state.push(v);
                             break 'select_fork state.save();
                         }
                     }
                 }
-                (OnFork::IterateInput, Some(_e)) => {
-                    env.push_fork(state, OnFork::IterateInput, state.pc);
+                (OnFork::IterateContext, Some(_e)) => {
+                    env.push_fork(state, OnFork::IterateContext, state.pc);
                     return Some(Err(err.take().unwrap()));
                 }
                 (_, Some(_)) => continue 'select_fork,


Test patch:
diff --git a/tests/common/mod.rs b/tests/common/mod.rs
index 570c7d0..48f218d 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, InputError, Value};
+use xq::{module_loader::PreludeLoader, run_query, util::SharedIterator, InputError, Value};
 
 #[macro_export]
 macro_rules! test {
@@ -12,13 +12,15 @@ macro_rules! test {
 }
 
 pub(crate) fn run_test(query: &str, input: &str, output: &str) -> Result<(), Box<dyn Error>> {
-    let input = serde_json::de::Deserializer::from_str(input)
+    let input: SharedIterator<_> = serde_json::de::Deserializer::from_str(input)
         .into_iter::<Value>()
-        .map(|r| r.map_err(InputError::new));
+        .map(|r| r.map_err(InputError::new))
+        .into_iter()
+        .into();
     let expected: Vec<_> = serde_json::de::Deserializer::from_str(output)
         .into_iter::<Value>()
         .collect::<Result<_, serde_json::Error>>()?;
-    let output = run_query(query, input.into_iter(), &PreludeLoader())?
+    let output = run_query(query, input.clone(), input, &PreludeLoader())?
         .collect::<Result<Vec<Value>, _>>()?;
     if expected != output {
         eprintln!("{:?} {:?}", expected, output);