Add lazy closure evaluation to default (#14160) (#15654)

# Description

This PR adds lazy closure evaluation to the `default` command (closes
#14160).

- For non-closure values and without providing a column name, `default`
acts the same as before
- The user can now provide multiple column names to populate if empty
- If the user provides a column name, the input must be a record or
list, otherwise an error is created.
- The user can now provide a closure as a default value
  - This closure is run without any arguments or input
  - The closure is never evaluated if the value isn't needed
- Even when column names are supplied, the closure is only run once (and
cached to prevent re-calling it)

For example:

```nushell
> default { 1 + 2 } # => 3
> null | default 3 a   # => previously `null`, now errors
> 1 | default { sleep 5sec; 3 } # => `1`, without waiting 5 seconds

> let optional_var = null; $optional_var | default { input 'Enter value: ' } # => Returns user input
> 5 | default { input 'Enter value: ' } # => `5`, without prompting user

> ls | default { sleep 5sec; 'N/A' } name # => No-op since `name` column is never empty
> ls | default { sleep 5sec; 'N/A' } foo bar # => creates columns `foo` and `bar`; only takes 5 seconds since closure result is cached

# Old behavior is the same
> [] | default 'foo' # => []
> [] | default --empty 'foo' # => 'foo'
> default 5 # => 5
```

# User-Facing Changes

- Users can add default values to multiple columns now.
- Users can now use closures as the default value passed to `default`.
- To return a closure, the user must wrap the closure they want to
return inside another closure, which will be run (`default { $my_closure
}`).

# Tests + Formatting

All tests pass.

# After Submitting

---------

Co-authored-by: 132ikl <132@ikl.sh>

crates/nu-command/src/filters/default.rs

@@ -1,4 +1,4 @@
-use nu_engine::command_prelude::*;
+use nu_engine::{ClosureEval, command_prelude::*};
 use nu_protocol::{ListStream, Signals};
 
 #[derive(Clone)]
@@ -19,7 +19,7 @@ impl Command for Default {
                 SyntaxShape::Any,
                 "The value to use as a default.",
             )
-            .optional(
+            .rest(
                 "column name",
                 SyntaxShape::String,
                 "The name of the column.",
@@ -43,8 +43,18 @@ impl Command for Default {
         call: &Call,
         input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
+        let default_value: Value = call.req(engine_state, stack, 0)?;
+        let columns: Vec<String> = call.rest(engine_state, stack, 1)?;
         let empty = call.has_flag(engine_state, stack, "empty")?;
-        default(engine_state, stack, call, input, empty)
+        default(
+            engine_state,
+            stack,
+            call,
+            input,
+            default_value,
+            empty,
+            columns,
+        )
     }
 
     fn examples(&self) -> Vec<Example> {
@@ -100,6 +110,25 @@ impl Command for Default {
                     }),
                 ])),
             },
+            Example {
+                description: r#"Generate a default value from a closure"#,
+                example: "null | default { 1 + 2 }",
+                result: Some(Value::test_int(3)),
+            },
+            Example {
+                description: r#"Fill missing column values based on other columns"#,
+                example: r#"[{a:1 b:2} {b:1}] | upsert a {|rc| default { $rc.b + 1 } }"#,
+                result: Some(Value::test_list(vec![
+                    Value::test_record(record! {
+                        "a" => Value::test_int(1),
+                        "b" => Value::test_int(2),
+                    }),
+                    Value::test_record(record! {
+                        "a" => Value::test_int(2),
+                        "b" => Value::test_int(1),
+                    }),
+                ])),
+            },
         ]
     }
 }
@@ -109,43 +138,68 @@ fn default(
     stack: &mut Stack,
     call: &Call,
     input: PipelineData,
+    default_value: Value,
     default_when_empty: bool,
+    columns: Vec<String>,
 ) -> Result<PipelineData, ShellError> {
+    let input_span = input.span().unwrap_or(call.head);
+    let mut default_value = DefaultValue::new(engine_state, stack, default_value);
     let metadata = input.metadata();
-    let value: Value = call.req(engine_state, stack, 0)?;
-    let column: Option<Spanned<String>> = call.opt(engine_state, stack, 1)?;
-
-    if let Some(column) = column {
-        input
-            .map(
-                move |mut item| match item {
-                    Value::Record {
-                        val: ref mut record,
-                        ..
-                    } => {
-                        let record = record.to_mut();
-                        if let Some(val) = record.get_mut(&column.item) {
-                            if matches!(val, Value::Nothing { .. })
-                                || (default_when_empty && val.is_empty())
-                            {
-                                *val = value.clone();
-                            }
-                        } else {
-                            record.push(column.item.clone(), value.clone());
-                        }
 
+    // If user supplies columns, check if input is a record or list of records
+    // and set the default value for the specified record columns
+    if !columns.is_empty() {
+        if matches!(input, PipelineData::Value(Value::Record { .. }, _)) {
+            let record = input.into_value(input_span)?.into_record()?;
+            fill_record(
+                record,
+                input_span,
+                &mut default_value,
+                columns.as_slice(),
+                default_when_empty,
+            )
+            .map(|x| x.into_pipeline_data_with_metadata(metadata))
+        } else if matches!(
+            input,
+            PipelineData::ListStream(..) | PipelineData::Value(Value::List { .. }, _)
+        ) {
+            // Potential enhancement: add another branch for Value::List,
+            // and collect the iterator into a Result<Value::List, ShellError>
+            // so we can preemptively return an error for collected lists
+            let head = call.head;
+            Ok(input
+                .into_iter()
+                .map(move |item| {
+                    let span = item.span();
+                    if let Value::Record { val, .. } = item {
+                        fill_record(
+                            val.into_owned(),
+                            span,
+                            &mut default_value,
+                            columns.as_slice(),
+                            default_when_empty,
+                        )
+                        .unwrap_or_else(|err| Value::error(err, head))
+                    } else {
                         item
                     }
-                    _ => item,
+                })
-                },
+                .into_pipeline_data_with_metadata(head, engine_state.signals().clone(), metadata))
-                engine_state.signals(),
+        // If columns are given, but input does not use columns, return an error
-            )
+        } else {
-            .map(|x| x.set_metadata(metadata))
+            Err(ShellError::PipelineMismatch {
+                exp_input_type: "record, table".to_string(),
+                dst_span: input_span,
+                src_span: input_span,
+            })
+        }
+    // Otherwise, if no column name is given, check if value is null
+    // or an empty string, list, or record when --empty is passed
     } else if input.is_nothing()
         || (default_when_empty
             && matches!(input, PipelineData::Value(ref value, _) if value.is_empty()))
     {
-        Ok(value.into_pipeline_data())
+        default_value.pipeline_data()
     } else if default_when_empty && matches!(input, PipelineData::ListStream(..)) {
         let PipelineData::ListStream(ls, metadata) = input else {
             unreachable!()
@@ -153,26 +207,86 @@ fn default(
         let span = ls.span();
         let mut stream = ls.into_inner().peekable();
         if stream.peek().is_none() {
-            return Ok(value.into_pipeline_data());
+            return default_value.pipeline_data();
         }
 
         // stream's internal state already preserves the original signals config, so if this
         // Signals::empty list stream gets interrupted it will be caught by the underlying iterator
         let ls = ListStream::new(stream, span, Signals::empty());
         Ok(PipelineData::ListStream(ls, metadata))
+    // Otherwise, return the input as is
     } else {
         Ok(input)
     }
 }
 
+/// A wrapper around the default value to handle closures and caching values
+enum DefaultValue {
+    Uncalculated(Spanned<ClosureEval>),
+    Calculated(Value),
+}
+
+impl DefaultValue {
+    fn new(engine_state: &EngineState, stack: &Stack, value: Value) -> Self {
+        let span = value.span();
+        match value {
+            Value::Closure { val, .. } => {
+                let closure_eval = ClosureEval::new(engine_state, stack, *val);
+                DefaultValue::Uncalculated(closure_eval.into_spanned(span))
+            }
+            _ => DefaultValue::Calculated(value),
+        }
+    }
+
+    fn value(&mut self) -> Result<Value, ShellError> {
+        match self {
+            DefaultValue::Uncalculated(closure) => {
+                let value = closure
+                    .item
+                    .run_with_input(PipelineData::Empty)?
+                    .into_value(closure.span)?;
+                *self = DefaultValue::Calculated(value.clone());
+                Ok(value)
+            }
+            DefaultValue::Calculated(value) => Ok(value.clone()),
+        }
+    }
+
+    fn pipeline_data(&mut self) -> Result<PipelineData, ShellError> {
+        self.value().map(|x| x.into_pipeline_data())
+    }
+}
+
+/// Given a record, fill missing columns with a default value
+fn fill_record(
+    mut record: Record,
+    span: Span,
+    default_value: &mut DefaultValue,
+    columns: &[String],
+    empty: bool,
+) -> Result<Value, ShellError> {
+    for col in columns {
+        if let Some(val) = record.get_mut(col) {
+            if matches!(val, Value::Nothing { .. }) || (empty && val.is_empty()) {
+                *val = default_value.value()?;
+            }
+        } else {
+            record.push(col.clone(), default_value.value()?);
+        }
+    }
+    Ok(Value::record(record, span))
+}
+
 #[cfg(test)]
 mod test {
+    use crate::Upsert;
+
     use super::*;
 
     #[test]
     fn test_examples() {
-        use crate::test_examples;
+        use crate::test_examples_with_commands;
 
-        test_examples(Default {})
+        test_examples_with_commands(Default {}, &[&Upsert]);
     }
 }

crates/nu-command/tests/commands/default.rs

@@ -150,3 +150,97 @@ fn do_not_replace_non_empty_list_stream() {
         assert_eq!(actual.out, "2");
     })
 }
+
+#[test]
+fn closure_eval_simple() {
+    let actual = nu!(r#"null | default { 1 }"#);
+    assert_eq!(actual.out, "1");
+}
+
+#[test]
+fn closure_eval_complex() {
+    let actual = nu!(r#"null | default { seq 1 5 | math sum }"#);
+    assert_eq!(actual.out, "15");
+}
+
+#[test]
+fn closure_eval_is_lazy() {
+    let actual = nu!(r#"1 | default { error make -u {msg: foo} }"#);
+    assert_eq!(actual.out, "1");
+}
+
+#[test]
+fn column_closure_eval_is_lazy() {
+    let actual = nu!(r#"{a: 1} | default { error make -u {msg: foo} } a | get a"#);
+    assert_eq!(actual.out, "1");
+}
+
+#[test]
+fn closure_eval_replace_empty_string() {
+    let actual = nu!(r#"'' | default --empty { 1 }"#);
+    assert_eq!(actual.out, "1");
+}
+
+#[test]
+fn closure_eval_do_not_replace_empty_string() {
+    let actual = nu!(r#"'' | default { 1 }"#);
+    assert_eq!(actual.out, "");
+}
+
+#[test]
+fn closure_eval_replace_empty_list() {
+    let actual = nu!(r#"[] | default --empty { 1 }"#);
+    assert_eq!(actual.out, "1");
+}
+
+#[test]
+fn closure_eval_do_not_replace_empty_list() {
+    let actual = nu!(r#"[] | default { 1 } | length"#);
+    assert_eq!(actual.out, "0");
+}
+
+#[test]
+fn closure_eval_replace_empty_record() {
+    let actual = nu!(r#"{} | default --empty { 1 }"#);
+    assert_eq!(actual.out, "1");
+}
+
+#[test]
+fn closure_eval_do_not_replace_empty_record() {
+    let actual = nu!(r#"{} | default { 1 } | columns | length"#);
+    assert_eq!(actual.out, "0");
+}
+
+#[test]
+fn closure_eval_add_missing_column_record() {
+    let actual = nu!(r#"
+        {a: 1} | default { 2 } b | get b
+    "#);
+    assert_eq!(actual.out, "2");
+}
+
+#[test]
+fn closure_eval_add_missing_column_table() {
+    let actual = nu!(r#"
+        [{a: 1, b: 2}, {b: 4}] | default { 3 } a | get a | to json -r
+    "#);
+    assert_eq!(actual.out, "[1,3]");
+}
+
+#[test]
+fn closure_eval_replace_empty_column() {
+    let actual = nu!(r#"{a: ''} | default -e { 1 } a | get a"#);
+    assert_eq!(actual.out, "1");
+}
+
+#[test]
+fn replace_multiple_columns() {
+    let actual = nu!(r#"{a: ''} | default -e 1 a b | values | to json -r"#);
+    assert_eq!(actual.out, "[1,1]");
+}
+
+#[test]
+fn return_closure_value() {
+    let actual = nu!(r#"null | default { {||} }"#);
+    assert!(actual.out.starts_with("closure"));
+}