more instructions for read/write variables, cell path, env var

crates/nu-engine/src/compile/mod.rs

@@ -1,11 +1,11 @@
 use nu_protocol::{
     ast::{
-        Argument, Block, Call, CellPath, Expr, Expression, Operator, Pipeline, PipelineRedirection,
+        Argument, Block, Call, CellPath, Expr, Expression, Operator, PathMember, Pipeline,
-        RedirectionSource, RedirectionTarget,
+        PipelineRedirection, RedirectionSource, RedirectionTarget,
     },
     engine::StateWorkingSet,
     ir::{Instruction, IrBlock, Literal, RedirectMode},
-    IntoSpanned, OutDest, RegId, ShellError, Span, Spanned,
+    IntoSpanned, OutDest, RegId, ShellError, Span, Spanned, ENV_VARIABLE_ID,
 };
 
 const BLOCK_INPUT: RegId = RegId(0);
@@ -264,7 +264,13 @@ fn compile_expression(
         Expr::Float(f) => lit(builder, Literal::Float(*f)),
         Expr::Binary(bin) => lit(builder, Literal::Binary(bin.as_slice().into())),
         Expr::Range(_) => Err(CompileError::Todo("Range")),
-        Expr::Var(_) => Err(CompileError::Todo("Var")),
+        Expr::Var(var_id) => builder.push(
+            Instruction::LoadVariable {
+                dst: out_reg,
+                var_id: *var_id,
+            }
+            .into_spanned(expr.span),
+        ),
         Expr::VarDecl(_) => Err(CompileError::Todo("VarDecl")),
         Expr::Call(call) => {
             // Ensure that out_reg contains the input value, because a call only uses one register
@@ -332,38 +338,42 @@ fn compile_expression(
         Expr::RawString(rs) => lit(builder, Literal::RawString(rs.as_str().into())),
         Expr::CellPath(path) => lit(builder, Literal::CellPath(Box::new(path.clone()))),
         Expr::FullCellPath(full_cell_path) => {
-            compile_expression(
+            if matches!(full_cell_path.head.expr, Expr::Var(ENV_VARIABLE_ID)) {
-                working_set,
+                compile_load_env(builder, expr.span, &full_cell_path.tail, out_reg)
-                builder,
+            } else {
-                &full_cell_path.head,
+                compile_expression(
-                RedirectModes::capture_out(expr.span),
+                    working_set,
-                in_reg,
+                    builder,
-                out_reg,
+                    &full_cell_path.head,
-            )?;
+                    RedirectModes::capture_out(expr.span),
-            // Only do the follow if this is actually needed
+                    in_reg,
-            if !full_cell_path.tail.is_empty() {
+                    out_reg,
-                let cell_path_reg = builder.literal(
-                    Literal::CellPath(Box::new(CellPath {
-                        members: full_cell_path.tail.clone(),
-                    }))
-                    .into_spanned(expr.span),
-                )?;
-                builder.push(
-                    Instruction::FollowCellPath {
-                        src_dst: out_reg,
-                        path: cell_path_reg,
-                    }
-                    .into_spanned(expr.span),
                 )?;
+                // Only do the follow if this is actually needed
+                if !full_cell_path.tail.is_empty() {
+                    let cell_path_reg = builder.literal(
+                        Literal::CellPath(Box::new(CellPath {
+                            members: full_cell_path.tail.clone(),
+                        }))
+                        .into_spanned(expr.span),
+                    )?;
+                    builder.push(
+                        Instruction::FollowCellPath {
+                            src_dst: out_reg,
+                            path: cell_path_reg,
+                        }
+                        .into_spanned(expr.span),
+                    )?;
+                }
+                Ok(())
             }
-            Ok(())
         }
         Expr::ImportPattern(_) => Err(CompileError::Todo("ImportPattern")),
         Expr::Overlay(_) => Err(CompileError::Todo("Overlay")),
         Expr::Signature(_) => Err(CompileError::Todo("Signature")),
         Expr::StringInterpolation(_) => Err(CompileError::Todo("StringInterpolation")),
-        Expr::Nothing => Err(CompileError::Todo("Nothing ")),
+        Expr::Nothing => Err(CompileError::Todo("Nothing")),
-        Expr::Garbage => Err(CompileError::Todo("Garbage ")),
+        Expr::Garbage => Err(CompileError::Todo("Garbage")),
     }
 }
 
@@ -519,6 +529,53 @@ fn compile_binary_op(
     Ok(())
 }
 
+fn compile_load_env(
+    builder: &mut BlockBuilder,
+    span: Span,
+    path: &[PathMember],
+    out_reg: RegId,
+) -> Result<(), CompileError> {
+    if path.is_empty() {
+        builder.push(
+            Instruction::LoadVariable {
+                dst: out_reg,
+                var_id: ENV_VARIABLE_ID,
+            }
+            .into_spanned(span),
+        )
+    } else {
+        let (key, optional) = match &path[0] {
+            PathMember::String { val, optional, .. } => (val.as_str().into(), *optional),
+            PathMember::Int { span, .. } => return Err(CompileError::AccessEnvByInt(*span)),
+        };
+        let tail = &path[1..];
+
+        if optional {
+            builder.push(Instruction::LoadEnvOpt { dst: out_reg, key }.into_spanned(span))?;
+        } else {
+            builder.push(Instruction::LoadEnv { dst: out_reg, key }.into_spanned(span))?;
+        }
+
+        if !tail.is_empty() {
+            let path = builder.literal(
+                Literal::CellPath(Box::new(CellPath {
+                    members: tail.to_vec(),
+                }))
+                .into_spanned(span),
+            )?;
+            builder.push(
+                Instruction::FollowCellPath {
+                    src_dst: out_reg,
+                    path,
+                }
+                .into_spanned(span),
+            )?;
+        }
+
+        Ok(())
+    }
+}
+
 /// An internal compiler error, generally means a Nushell bug rather than an issue with user error
 /// since parsing and typechecking has already passed.
 #[derive(Debug)]
@@ -528,11 +585,12 @@ enum CompileError {
     InvalidRedirectMode,
     Garbage,
     UnsupportedOperatorExpression,
+    AccessEnvByInt(Span),
     Todo(&'static str),
 }
 
 impl CompileError {
-    fn to_shell_error(self, span: Option<Span>) -> ShellError {
+    fn to_shell_error(self, mut span: Option<Span>) -> ShellError {
         let ice = "internal compiler error: ";
         let message = match self {
             CompileError::RegisterOverflow => format!("{ice}register overflow"),
@@ -548,6 +606,10 @@ impl CompileError {
             CompileError::UnsupportedOperatorExpression => {
                 format!("{ice}unsupported operator expression")
             }
+            CompileError::AccessEnvByInt(local_span) => {
+                span = Some(local_span);
+                format!("{ice}attempted access of $env by integer path")
+            }
             CompileError::Todo(msg) => {
                 format!("{ice}TODO: {msg}")
             }
@@ -631,8 +693,8 @@ impl BlockBuilder {
         }
     }
 
-    /// Insert an instruction into the block, automatically freeing any registers consumed by the
+    /// Insert an instruction into the block, automatically marking any registers populated by
-    /// instruction.
+    /// the instruction, and freeing any registers consumed by the instruction.
     fn push(&mut self, instruction: Spanned<Instruction>) -> Result<(), CompileError> {
         match &instruction.item {
             Instruction::LoadLiteral { dst, lit: _ } => self.mark_register(*dst)?,
@@ -643,6 +705,11 @@ impl BlockBuilder {
             Instruction::Clone { dst, src: _ } => self.mark_register(*dst)?,
             Instruction::Collect { src_dst: _ } => (),
             Instruction::Drain { src } => self.free_register(*src)?,
+            Instruction::LoadVariable { dst, var_id: _ } => self.mark_register(*dst)?,
+            Instruction::StoreVariable { var_id: _, src } => self.free_register(*src)?,
+            Instruction::LoadEnv { dst, key: _ } => self.mark_register(*dst)?,
+            Instruction::LoadEnvOpt { dst, key: _ } => self.mark_register(*dst)?,
+            Instruction::StoreEnv { key: _, src } => self.free_register(*src)?,
             Instruction::PushPositional { src } => self.free_register(*src)?,
             Instruction::AppendRest { src } => self.free_register(*src)?,
             Instruction::PushFlag { name: _ } => (),
@@ -661,6 +728,18 @@ impl BlockBuilder {
                 rhs,
             } => self.free_register(*rhs)?,
             Instruction::FollowCellPath { src_dst: _, path } => self.free_register(*path)?,
+            Instruction::CloneCellPath { dst, src: _, path } => {
+                self.mark_register(*dst)?;
+                self.free_register(*path)?;
+            }
+            Instruction::UpsertCellPath {
+                src_dst: _,
+                path,
+                new_value,
+            } => {
+                self.free_register(*path)?;
+                self.free_register(*new_value)?;
+            }
             Instruction::Jump { index: _ } => (),
             Instruction::BranchIf { cond, index: _ } => self.free_register(*cond)?,
             Instruction::Return { src } => self.free_register(*src)?,

crates/nu-protocol/src/ir/display.rs

@@ -1,6 +1,6 @@
 use std::fmt;
 
-use crate::{engine::EngineState, DeclId};
+use crate::{engine::EngineState, DeclId, VarId};
 
 use super::{Instruction, IrBlock, RedirectMode};
 
@@ -60,6 +60,23 @@ impl<'a> fmt::Display for FmtInstruction<'a> {
             Instruction::Drain { src } => {
                 write!(f, "{:WIDTH$} {src}", "drain")
             }
+            Instruction::LoadVariable { dst, var_id } => {
+                let var = FmtVar::new(self.engine_state, *var_id);
+                write!(f, "{:WIDTH$} {dst}, {var}", "load-variable")
+            }
+            Instruction::StoreVariable { var_id, src } => {
+                let var = FmtVar::new(self.engine_state, *var_id);
+                write!(f, "{:WIDTH$} {var}, {src}", "store-variable")
+            }
+            Instruction::LoadEnv { dst, key } => {
+                write!(f, "{:WIDTH$} {dst}, {key:?}", "load-env")
+            }
+            Instruction::LoadEnvOpt { dst, key } => {
+                write!(f, "{:WIDTH$} {dst}, {key:?}", "load-env-opt")
+            }
+            Instruction::StoreEnv { key, src } => {
+                write!(f, "{:WIDTH$} {key:?}, {src}", "store-env")
+            }
             Instruction::PushPositional { src } => {
                 write!(f, "{:WIDTH$} {src}", "push-positional")
             }
@@ -88,6 +105,20 @@ impl<'a> fmt::Display for FmtInstruction<'a> {
             Instruction::FollowCellPath { src_dst, path } => {
                 write!(f, "{:WIDTH$} {src_dst}, {path}", "follow-cell-path")
             }
+            Instruction::CloneCellPath { dst, src, path } => {
+                write!(f, "{:WIDTH$} {dst}, {src}, {path}", "clone-cell-path")
+            }
+            Instruction::UpsertCellPath {
+                src_dst,
+                path,
+                new_value,
+            } => {
+                write!(
+                    f,
+                    "{:WIDTH$} {src_dst}, {path}, {new_value}",
+                    "upsert-cell-path"
+                )
+            }
             Instruction::Jump { index } => {
                 write!(f, "{:WIDTH$} {index}", "jump")
             }
@@ -115,6 +146,30 @@ impl fmt::Display for FmtDecl<'_> {
     }
 }
 
+struct FmtVar<'a>(DeclId, Option<&'a str>);
+
+impl<'a> FmtVar<'a> {
+    fn new(engine_state: &'a EngineState, var_id: VarId) -> Self {
+        // Search for the name of the variable
+        let name: Option<&str> = engine_state
+            .active_overlays(&[])
+            .flat_map(|overlay| overlay.vars.iter())
+            .find(|(_, v)| **v == var_id)
+            .map(|(k, _)| std::str::from_utf8(k).unwrap_or("<utf-8 error>"));
+        FmtVar(var_id, name)
+    }
+}
+
+impl fmt::Display for FmtVar<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if let Some(name) = self.1 {
+            write!(f, "var {} {:?}", self.0, name)
+        } else {
+            write!(f, "var {}", self.0)
+        }
+    }
+}
+
 impl std::fmt::Display for RedirectMode {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {

crates/nu-protocol/src/ir/mod.rs

@@ -1,7 +1,7 @@
 use crate::{
     ast::{CellPath, Operator},
     engine::EngineState,
-    BlockId, DeclId, RegId, Span,
+    BlockId, DeclId, RegId, Span, VarId,
 };
 
 use serde::{Deserialize, Serialize};
@@ -39,6 +39,17 @@ pub enum Instruction {
     Collect { src_dst: RegId },
     /// Drain the value/stream in a register and discard (e.g. semicolon)
     Drain { src: RegId },
+    /// Load the value of a variable into the `dst` register
+    LoadVariable { dst: RegId, var_id: VarId },
+    /// Store the value of a variable from the `src` register
+    StoreVariable { var_id: VarId, src: RegId },
+    /// Load the value of an environment variable into the `dst` register
+    LoadEnv { dst: RegId, key: Box<str> },
+    /// Load the value of an environment variable into the `dst` register, or `Nothing` if it
+    /// doesn't exist
+    LoadEnvOpt { dst: RegId, key: Box<str> },
+    /// Store the value of an environment variable from the `src` register
+    StoreEnv { key: Box<str>, src: RegId },
     /// Add a positional arg to the next call
     PushPositional { src: RegId },
     /// Add a list of args to the next call (spread/rest)
@@ -61,8 +72,19 @@ pub enum Instruction {
         op: Operator,
         rhs: RegId,
     },
-    /// Follow a cell path on the `path`
+    /// Follow a cell path on the value in `src_dst`, storing the result back to `src_dst`
     FollowCellPath { src_dst: RegId, path: RegId },
+    /// Clone the value at a cell path in `src`, storing the result to `dst`. The original value
+    /// remains in `src`. Must be a collected value.
+    CloneCellPath { dst: RegId, src: RegId, path: RegId },
+    /// Update/insert a cell path to `new_value` on the value in `src_dst`, storing the modified
+    /// value back to `src_dst`
+    UpsertCellPath {
+        src_dst: RegId,
+        path: RegId,
+        new_value: RegId,
+    },
+    /// Update a cell path
     /// Jump to an offset in this block
     Jump { index: usize },
     /// Branch to an offset in this block if the value of the `cond` register is a true boolean,