Allow parse-time evaluation of calls, pipelines and subexpressions (#9499)

Co-authored-by: Antoine Stevan <44101798+amtoine@users.noreply.github.com>

crates/nu-protocol/src/engine/command.rs

@@ -2,7 +2,7 @@ use std::path::PathBuf;
 
 use crate::{ast::Call, Alias, BlockId, Example, PipelineData, ShellError, Signature};
 
-use super::{EngineState, Stack};
+use super::{EngineState, Stack, StateWorkingSet};
 
 #[derive(Debug)]
 pub enum CommandType {
@@ -34,6 +34,19 @@ pub trait Command: Send + Sync + CommandClone {
         input: PipelineData,
     ) -> Result<PipelineData, ShellError>;
 
+    /// Used by the parser to run command at parse time
+    ///
+    /// If a command has `is_const()` set to true, it must also implement this method.
+    #[allow(unused_variables)]
+    fn run_const(
+        &self,
+        working_set: &StateWorkingSet,
+        call: &Call,
+        input: PipelineData,
+    ) -> Result<PipelineData, ShellError> {
+        Err(ShellError::MissingConstEvalImpl { span: call.head })
+    }
+
     fn examples(&self) -> Vec<Example> {
         Vec::new()
     }
@@ -83,6 +96,11 @@ pub trait Command: Send + Sync + CommandClone {
         None
     }
 
+    // Whether can run in const evaluation in the parser
+    fn is_const(&self) -> bool {
+        false
+    }
+
     // If command is a block i.e. def blah [] { }, get the block id
     fn get_block_id(&self) -> Option<BlockId> {
         None

crates/nu-protocol/src/engine/engine_state.rs

@@ -19,7 +19,7 @@ use std::sync::{
     Arc, Mutex,
 };
 
-static PWD_ENV: &str = "PWD";
+pub static PWD_ENV: &str = "PWD";
 
 /// Organizes usage messages for various primitives
 #[derive(Debug, Clone)]
@@ -1090,6 +1090,10 @@ impl<'a> StateWorkingSet<'a> {
         }
     }
 
+    pub fn permanent(&self) -> &EngineState {
+        self.permanent_state
+    }
+
     pub fn error(&mut self, parse_error: ParseError) {
         self.parse_errors.push(parse_error)
     }

crates/nu-protocol/src/eval_const.rs

@@ -0,0 +1,196 @@
+use crate::{
+    ast::{Block, Call, Expr, Expression, PipelineElement},
+    engine::StateWorkingSet,
+    PipelineData, Record, ShellError, Span, Value,
+};
+
+fn eval_const_call(
+    working_set: &StateWorkingSet,
+    call: &Call,
+    input: PipelineData,
+) -> Result<PipelineData, ShellError> {
+    let decl = working_set.get_decl(call.decl_id);
+
+    if !decl.is_const() {
+        return Err(ShellError::NotAConstCommand(call.head));
+    }
+
+    if !decl.is_known_external() && call.named_iter().any(|(flag, _, _)| flag.item == "help") {
+        // It would require re-implementing get_full_help() for const evaluation. Assuming that
+        // getting help messages at parse-time is rare enough, we can simply disallow it.
+        return Err(ShellError::NotAConstHelp(call.head));
+    }
+
+    decl.run_const(working_set, call, input)
+}
+
+fn eval_const_subexpression(
+    working_set: &StateWorkingSet,
+    expr: &Expression,
+    block: &Block,
+    mut input: PipelineData,
+) -> Result<PipelineData, ShellError> {
+    for pipeline in block.pipelines.iter() {
+        for element in pipeline.elements.iter() {
+            let PipelineElement::Expression(_, expr) = element else {
+                return Err(ShellError::NotAConstant(expr.span));
+            };
+
+            input = eval_constant_with_input(working_set, expr, input)?
+        }
+    }
+
+    Ok(input)
+}
+
+fn eval_constant_with_input(
+    working_set: &StateWorkingSet,
+    expr: &Expression,
+    input: PipelineData,
+) -> Result<PipelineData, ShellError> {
+    match &expr.expr {
+        Expr::Call(call) => eval_const_call(working_set, call, input),
+        Expr::Subexpression(block_id) => {
+            let block = working_set.get_block(*block_id);
+            eval_const_subexpression(working_set, expr, block, input)
+        }
+        _ => eval_constant(working_set, expr).map(|v| PipelineData::Value(v, None)),
+    }
+}
+
+/// Evaluate a constant value at parse time
+///
+/// Based off eval_expression() in the engine
+pub fn eval_constant(
+    working_set: &StateWorkingSet,
+    expr: &Expression,
+) -> Result<Value, ShellError> {
+    match &expr.expr {
+        Expr::Bool(b) => Ok(Value::bool(*b, expr.span)),
+        Expr::Int(i) => Ok(Value::int(*i, expr.span)),
+        Expr::Float(f) => Ok(Value::float(*f, expr.span)),
+        Expr::Binary(b) => Ok(Value::Binary {
+            val: b.clone(),
+            span: expr.span,
+        }),
+        Expr::Filepath(path) => Ok(Value::String {
+            val: path.clone(),
+            span: expr.span,
+        }),
+        Expr::Var(var_id) => match working_set.get_variable(*var_id).const_val.as_ref() {
+            Some(val) => Ok(val.clone()),
+            None => Err(ShellError::NotAConstant(expr.span)),
+        },
+        Expr::CellPath(cell_path) => Ok(Value::CellPath {
+            val: cell_path.clone(),
+            span: expr.span,
+        }),
+        Expr::FullCellPath(cell_path) => {
+            let value = eval_constant(working_set, &cell_path.head)?;
+
+            match value.follow_cell_path(&cell_path.tail, false) {
+                Ok(val) => Ok(val),
+                // TODO: Better error conversion
+                Err(shell_error) => Err(ShellError::GenericError(
+                    "Error when following cell path".to_string(),
+                    format!("{shell_error:?}"),
+                    Some(expr.span),
+                    None,
+                    vec![],
+                )),
+            }
+        }
+        Expr::DateTime(dt) => Ok(Value::Date {
+            val: *dt,
+            span: expr.span,
+        }),
+        Expr::List(x) => {
+            let mut output = vec![];
+            for expr in x {
+                output.push(eval_constant(working_set, expr)?);
+            }
+            Ok(Value::List {
+                vals: output,
+                span: expr.span,
+            })
+        }
+        Expr::Record(fields) => {
+            let mut record = Record::new();
+            for (col, val) in fields {
+                // avoid duplicate cols.
+                let col_name = value_as_string(eval_constant(working_set, col)?, expr.span)?;
+                let pos = record.cols.iter().position(|c| c == &col_name);
+                match pos {
+                    Some(index) => {
+                        record.vals[index] = eval_constant(working_set, val)?;
+                    }
+                    None => {
+                        record.push(col_name, eval_constant(working_set, val)?);
+                    }
+                }
+            }
+
+            Ok(Value::record(record, expr.span))
+        }
+        Expr::Table(headers, vals) => {
+            let mut output_headers = vec![];
+            for expr in headers {
+                output_headers.push(value_as_string(
+                    eval_constant(working_set, expr)?,
+                    expr.span,
+                )?);
+            }
+
+            let mut output_rows = vec![];
+            for val in vals {
+                let mut row = vec![];
+                for expr in val {
+                    row.push(eval_constant(working_set, expr)?);
+                }
+                output_rows.push(Value::record(
+                    Record {
+                        cols: output_headers.clone(),
+                        vals: row,
+                    },
+                    expr.span,
+                ));
+            }
+            Ok(Value::List {
+                vals: output_rows,
+                span: expr.span,
+            })
+        }
+        Expr::Keyword(_, _, expr) => eval_constant(working_set, expr),
+        Expr::String(s) => Ok(Value::String {
+            val: s.clone(),
+            span: expr.span,
+        }),
+        Expr::Nothing => Ok(Value::Nothing { span: expr.span }),
+        Expr::ValueWithUnit(expr, unit) => {
+            if let Ok(Value::Int { val, .. }) = eval_constant(working_set, expr) {
+                unit.item.to_value(val, unit.span)
+            } else {
+                Err(ShellError::NotAConstant(expr.span))
+            }
+        }
+        Expr::Call(call) => {
+            Ok(eval_const_call(working_set, call, PipelineData::empty())?.into_value(expr.span))
+        }
+        Expr::Subexpression(block_id) => {
+            let block = working_set.get_block(*block_id);
+            Ok(
+                eval_const_subexpression(working_set, expr, block, PipelineData::empty())?
+                    .into_value(expr.span),
+            )
+        }
+        _ => Err(ShellError::NotAConstant(expr.span)),
+    }
+}
+
+/// Get the value as a string
+pub fn value_as_string(value: Value, span: Span) -> Result<String, ShellError> {
+    match value {
+        Value::String { val, .. } => Ok(val),
+        _ => Err(ShellError::NotAConstant(span)),
+    }
+}

crates/nu-protocol/src/lib.rs

@@ -4,6 +4,7 @@ pub mod cli_error;
 pub mod config;
 mod did_you_mean;
 pub mod engine;
+pub mod eval_const;
 mod example;
 mod exportable;
 mod id;

crates/nu-protocol/src/parse_error.rs

@@ -449,18 +449,6 @@ pub enum ParseError {
     #[diagnostic(code(nu::shell::error_reading_file))]
     ReadingFile(String, #[label("{0}")] Span),
 
-    /// Tried assigning non-constant value to a constant
-    ///
-    /// ## Resolution
-    ///
-    /// Only a subset of expressions are allowed to be assigned as a constant during parsing.
-    #[error("Not a constant.")]
-    #[diagnostic(
-        code(nu::parser::not_a_constant),
-        help("Only a subset of expressions are allowed constants during parsing. Try using the 'const' command or typing the value literally.")
-    )]
-    NotAConstant(#[label = "Value is not a parse-time constant"] Span),
-
     #[error("Invalid literal")] // <problem> in <entity>.
     #[diagnostic()]
     InvalidLiteral(String, String, #[label("{0} in {1}")] Span),
@@ -561,7 +549,6 @@ impl ParseError {
             ParseError::ShellOutErrRedirect(s) => *s,
             ParseError::UnknownOperator(_, _, s) => *s,
             ParseError::InvalidLiteral(_, _, s) => *s,
-            ParseError::NotAConstant(s) => *s,
             ParseError::LabeledErrorWithHelp { span: s, .. } => *s,
         }
     }

crates/nu-protocol/src/shell_error.rs

@@ -2,7 +2,7 @@ use miette::Diagnostic;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
-use crate::{ast::Operator, Span, Value};
+use crate::{ast::Operator, engine::StateWorkingSet, format_error, ParseError, Span, Value};
 
 /// The fundamental error type for the evaluation engine. These cases represent different kinds of errors
 /// the evaluator might face, along with helpful spans to label. An error renderer will take this error value
@@ -1074,6 +1074,72 @@ pub enum ShellError {
         #[label("not a boolean expression")]
         span: Span,
     },
+
+    /// An attempt to run a command marked for constant evaluation lacking the const. eval.
+    /// implementation.
+    ///
+    /// This is an internal Nushell error, please file an issue.
+    #[error("Missing const eval implementation")]
+    #[diagnostic(
+        code(nu::shell::missing_const_eval_implementation),
+        help(
+            "The command lacks an implementation for constant evaluation. \
+This is an internal Nushell error, please file an issue https://github.com/nushell/nushell/issues."
+        )
+    )]
+    MissingConstEvalImpl {
+        #[label("command lacks constant implementation")]
+        span: Span,
+    },
+
+    /// Tried assigning non-constant value to a constant
+    ///
+    /// ## Resolution
+    ///
+    /// Only a subset of expressions are allowed to be assigned as a constant during parsing.
+    #[error("Not a constant.")]
+    #[diagnostic(
+        code(nu::shell::not_a_constant),
+        help("Only a subset of expressions are allowed constants during parsing. Try using the 'const' command or typing the value literally.")
+    )]
+    NotAConstant(#[label = "Value is not a parse-time constant"] Span),
+
+    /// Tried running a command that is not const-compatible
+    ///
+    /// ## Resolution
+    ///
+    /// Only a subset of builtin commands, and custom commands built only from those commands, can
+    /// run at parse time.
+    #[error("Not a const command.")]
+    #[diagnostic(
+        code(nu::shell::not_a_const_command),
+        help("Only a subset of builtin commands, and custom commands built only from those commands, can run at parse time.")
+    )]
+    NotAConstCommand(#[label = "This command cannot run at parse time."] Span),
+
+    /// Tried getting a help message at parse time.
+    ///
+    /// ## Resolution
+    ///
+    /// Help messages are not supported at parse time.
+    #[error("Help message not a constant.")]
+    #[diagnostic(
+        code(nu::shell::not_a_const_help),
+        help("Help messages are currently not supported to be constants.")
+    )]
+    NotAConstHelp(#[label = "Cannot get help message at parse time."] Span),
+}
+
+// TODO: Implement as From trait
+impl ShellError {
+    pub fn wrap(self, working_set: &StateWorkingSet, span: Span) -> ParseError {
+        let msg = format_error(working_set, &self);
+        ParseError::LabeledError(
+            msg,
+            "Encountered error during parse-time evaluation".into(),
+            span,
+        )
+    }
 }
 
 impl From<std::io::Error> for ShellError {