shrub_core/

callstack.rs

//! Call stack representations.

use std::borrow::Cow;
use std::collections::{HashSet, VecDeque};
use std::sync::Arc;

use shrub_parser::ast::SourceLocation;

use crate::{functions, traps};

/// Encapsulates info regarding a script call.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ScriptCall {
    /// The type of script call.
    pub call_type: ScriptCallType,
    /// The source info for the script called.
    pub source_info: crate::SourceInfo,
}

impl ScriptCall {
    /// Returns the name of the script that was called.
    pub fn name(&self) -> Cow<'_, str> {
        self.source_info.source.as_str().into()
    }
}

/// The type of script call.
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum ScriptCallType {
    /// A script was sourced.
    Source,
    /// A script was executed.
    Run,
}

impl std::fmt::Display for ScriptCall {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.call_type {
            ScriptCallType::Source => write!(f, "source({})", self.source_info),
            ScriptCallType::Run => write!(f, "script({})", self.source_info),
        }
    }
}

/// Represents the type of a frame, indicating how it was invoked from
/// a different source context.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum FrameType {
    /// A script was called (sourced or executed).
    Script(ScriptCall),
    /// A function was called.
    Function(FunctionCall),
    /// A trap handler was invoked.
    TrapHandler(traps::TrapSignal),
    /// A string was eval'd.
    Eval,
    /// A command-line string (i.e., -c) was executed.
    CommandString,
    /// An interactive command session was started.
    InteractiveSession,
}

impl FrameType {
    /// Returns a name for the frame (i.e., script path or function name).
    pub fn name(&self) -> Cow<'_, str> {
        match self {
            Self::Script(call) => call.name(),
            Self::Function(call) => call.name(),
            Self::TrapHandler(_) => "trap".into(),
            Self::Eval => "eval".into(),
            Self::CommandString => "-c".into(),
            Self::InteractiveSession => "interactive".into(),
        }
    }

    /// Returns `true` if the frame is for a function call.
    pub const fn is_function(&self) -> bool {
        matches!(self, Self::Function(..))
    }

    /// Returns `true` if the frame is for a script call.
    pub const fn is_script(&self) -> bool {
        matches!(self, Self::Script(..))
    }

    /// Returns `true` if the frame is for a trap handler.
    pub const fn is_trap_handler(&self) -> bool {
        matches!(self, Self::TrapHandler(_))
    }

    /// Returns `true` if the frame is for an interactive session.
    pub const fn is_interactive_session(&self) -> bool {
        matches!(self, Self::InteractiveSession)
    }

    /// Returns `true` if the frame is for a command string being executed.
    pub const fn is_command_string(&self) -> bool {
        matches!(self, Self::CommandString)
    }

    /// Returns `true` if the frame is for a sourced script.
    pub const fn is_sourced_script(&self) -> bool {
        matches!(self, Self::Script(call) if matches!(call.call_type, ScriptCallType::Source))
    }

    /// Returns `true` if the frame is for a run script.
    pub const fn is_run_script(&self) -> bool {
        matches!(self, Self::Script(call) if matches!(call.call_type, ScriptCallType::Run))
    }
}

impl std::fmt::Display for FrameType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Script(call) => call.fmt(f),
            Self::Function(call) => call.fmt(f),
            Self::TrapHandler(_) => write!(f, "trap"),
            Self::Eval => write!(f, "eval"),
            Self::CommandString => write!(f, "-c"),
            Self::InteractiveSession => write!(f, "interactive"),
        }
    }
}

/// Describes the target of a function call.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct FunctionCall {
    /// The name of the function invoked.
    pub function_name: String,
    /// The invoked function.
    pub function: functions::Registration,
}

impl FunctionCall {
    /// Returns the name of the function that was called.
    pub fn name(&self) -> Cow<'_, str> {
        self.function_name.as_str().into()
    }
}

impl std::fmt::Display for FunctionCall {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "func({})", self.function_name)
    }
}

/// Represents a single frame in a `CallStack`.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Frame {
    /// The type of frame.
    pub frame_type: FrameType,
    /// The source information for the frame. The locations associated with AST
    /// nodes executed in this frame should be interpreted as being relative
    /// to this source info.
    pub source_info: crate::SourceInfo,
    /// The location of the entry point into this frame, within the frame of
    /// reference of `source_info`. May be `None` if the entry point is not
    /// known.
    pub entry: Option<Arc<crate::SourcePosition>>,
    /// Information about the currently executing location. For the topmost
    /// frame on the stack, this represents the current execution location.
    /// For older frames, this represents the site from which a control
    /// transfer was made to the next younger frame. May be `None` if the
    /// current location is not known. When present, it is relative to the
    /// frame of reference of `source_info`.
    pub current: Option<Arc<crate::SourcePosition>>,
    /// Positional arguments (not including $0). May not be present for all
    /// frames.
    pub args: Vec<String>,
    /// Optionally, indicates an additional line offset within the current
    /// source context.
    pub current_line_offset: usize,
}

impl Frame {
    /// Returns the adjusted source info for this frame, combining the
    /// frame's `source_info` and `current_line_offset`, if present.
    pub fn adjusted_source_info(&self) -> crate::SourceInfo {
        self.pos_as_source_info(None)
    }

    /// Returns the current position as a new `SourceInfo`, combining the
    /// frame's `source_info` and `current` position.
    pub fn current_pos_as_source_info(&self) -> crate::SourceInfo {
        self.pos_as_source_info(self.current.as_ref())
    }

    fn pos_as_source_info(&self, pos: Option<&Arc<crate::SourcePosition>>) -> crate::SourceInfo {
        let mut new_start = if let Some(existing_start) = &self.source_info.start {
            if let Some(current) = pos {
                Some(Arc::new(crate::SourcePosition {
                    index: existing_start.index + current.index,
                    line: existing_start.line + (current.line - 1),
                    column: if current.line <= 1 {
                        existing_start.column + (current.column - 1)
                    } else {
                        current.column
                    },
                }))
            } else {
                Some(existing_start.clone())
            }
        } else {
            pos.cloned()
        };

        if self.current_line_offset > 0 {
            new_start = if let Some(new_start) = new_start {
                let mut pos = (*new_start).clone();
                pos.line += self.current_line_offset;

                Some(Arc::new(pos))
            } else {
                Some(Arc::new(crate::SourcePosition {
                    index: 0,
                    line: self.current_line_offset + 1,
                    column: 1,
                }))
            };
        }

        crate::SourceInfo {
            source: self.source_info.source.clone(),
            start: new_start,
        }
    }

    /// Returns the current line number.
    pub fn current_line(&self) -> Option<usize> {
        let start_line = self.source_info.start.as_ref().map_or(1, |pos| pos.line);
        let current_line = self.current.as_ref().map(|pos| pos.line)?;

        Some(start_line.saturating_sub(1) + current_line + self.current_line_offset)
    }

    /// Returns the current line number, relative to the frame's entry.
    pub fn current_frame_relative_line(&self) -> Option<usize> {
        let current_line = self.current.as_ref().map(|pos| pos.line)?;
        let entry_line = self.entry.as_ref().map_or(1, |pos| pos.line);

        Some(current_line.saturating_sub(entry_line) + self.current_line_offset + 1)
    }
}

/// Options for formatting a call stack.
#[derive(Default)]
pub struct FormatOptions {
    /// Whether or not to show args.
    pub show_args: bool,
    /// Whether or not to show frame entry points.
    pub show_entry_points: bool,
}

/// Helper struct for formatting a call stack with custom options.
///
/// This struct implements `Display` and can be used to write a formatted
/// call stack to any type that implements `io::Write`.
pub struct FormatCallStack<'a> {
    stack: &'a CallStack,
    options: &'a FormatOptions,
}

impl<'a> FormatCallStack<'a> {
    /// Creates a new formatter for the given call stack with the specified
    /// options.
    ///
    /// # Arguments
    ///
    /// * `stack` - The call stack to format.
    /// * `options` - The formatting options to use.
    pub const fn new(stack: &'a CallStack, options: &'a FormatOptions) -> Self {
        Self { stack, options }
    }
}

impl std::fmt::Display for FormatCallStack<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.stack.fmt_with_options(f, self.options)
    }
}

/// Encapsulates a script call stack.
#[derive(Clone, Debug, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct CallStack {
    frames: VecDeque<Frame>,
    func_call_depth: usize,
    script_source_depth: usize,
    active_trap_signals: HashSet<traps::TrapSignal>,
    trap_delivery_suppress_count: usize,
}

impl CallStack {
    /// Creates a formatter for this call stack with the given options.
    ///
    /// # Arguments
    ///
    /// * `options` - The formatting options to use.
    pub const fn format<'a>(&'a self, options: &'a FormatOptions) -> FormatCallStack<'a> {
        FormatCallStack::new(self, options)
    }

    /// Formats the call stack with the given options.
    ///
    /// # Arguments
    ///
    /// * `f` - The formatter to write to.
    /// * `options` - The formatting options.
    fn fmt_with_options(
        &self,
        f: &mut std::fmt::Formatter<'_>,
        options: &FormatOptions,
    ) -> std::fmt::Result {
        if self.is_empty() {
            return Ok(());
        }

        color_print::cwriteln!(f, "<underline>Call stack (most recent first):</underline>")?;

        for (index, frame) in self.iter().enumerate() {
            let si = frame.current_pos_as_source_info();

            color_print::cwrite!(
                f,
                "   <dim>#{index}</dim><yellow>|</yellow> <strong>{}</strong>",
                si.source
            )?;

            if let Some(pos) = &si.start {
                color_print::cwrite!(f, ":<cyan>{}</cyan>,<cyan>{}</cyan>", pos.line, pos.column)?;
            }

            color_print::cwrite!(f, " (<dim>{}</dim>", frame.frame_type)?;

            if options.show_entry_points {
                if let Some(entry) = &frame.entry {
                    let entry_si = frame.pos_as_source_info(Some(entry));
                    if let Some(entry_start) = &entry_si.start {
                        color_print::cwrite!(
                            f,
                            " <dim>entered at {}:{}</dim>",
                            entry_si.source,
                            entry_start
                        )?;
                    }
                }
            }

            color_print::cwriteln!(f, ")")?;

            if !frame.args.is_empty() && options.show_args {
                for (i, arg) in frame.args.iter().enumerate() {
                    color_print::cwriteln!(
                        f,
                        "     <yellow>${}</yellow>: <blue>{}</blue>",
                        i + 1,
                        arg
                    )?;
                }
            }
        }

        Ok(())
    }
}

impl std::fmt::Display for CallStack {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.fmt_with_options(f, &FormatOptions::default())
    }
}

impl std::ops::Index<usize> for CallStack {
    type Output = Frame;

    fn index(&self, index: usize) -> &Self::Output {
        &self.frames[index]
    }
}

impl CallStack {
    /// Creates a new empty script call stack.
    pub fn new() -> Self {
        Self::default()
    }

    /// Removes the top from from the stack. If the stack is empty, does nothing
    /// and returns `None`; otherwise, returns the removed call frame.
    pub fn pop(&mut self) -> Option<Frame> {
        let frame = self.frames.pop_front()?;

        if frame.frame_type.is_function() {
            self.func_call_depth = self.func_call_depth.saturating_sub(1);
        }

        if frame.frame_type.is_sourced_script() {
            self.script_source_depth = self.script_source_depth.saturating_sub(1);
        }

        if let FrameType::TrapHandler(signal) = &frame.frame_type {
            self.active_trap_signals.remove(signal);
        }

        Some(frame)
    }

    /// Returns a reference to the current (topmost) call frame in the stack.
    /// Returns `None` if the stack is empty.
    pub fn current_frame(&self) -> Option<&Frame> {
        self.frames.front()
    }

    /// Returns the position in the current (topmost) call frame in the stack,
    /// expressed as a new `SourceInfo`. Note that this may not be identical
    /// to that frame's `SourceInfo` since it may include an offset representing
    /// the current execution position within that source.
    pub fn current_pos_as_source_info(&self) -> crate::SourceInfo {
        let Some(frame) = self.frames.front() else {
            return crate::SourceInfo::default();
        };

        frame.current_pos_as_source_info()
    }

    /// Updates the currently executing position in the top stack frame.
    pub fn set_current_pos(&mut self, position: Option<Arc<crate::SourcePosition>>) {
        if let Some(frame) = self.frames.front_mut() {
            frame.current = position;
        }
    }

    /// Increments the current line offset in the top stack frame by the given
    /// delta.
    ///
    /// # Arguments
    ///
    /// * `delta` - The number of lines to increment the current line offset by.
    pub(crate) fn increment_current_line_offset(&mut self, delta: usize) {
        let Some(frame) = self.frames.front_mut() else {
            return;
        };

        frame.current_line_offset += delta;
    }

    /// Pushes a new script call frame onto the stack.
    ///
    /// # Arguments
    ///
    /// * `call_type` - The type of script call (sourced or executed).
    /// * `source_info` - The source of the script.
    /// * `args` - The positional arguments for the script call.
    pub fn push_script(
        &mut self,
        call_type: ScriptCallType,
        source_info: &crate::SourceInfo,
        args: impl IntoIterator<Item = String>,
    ) {
        self.frames.push_front(Frame {
            frame_type: FrameType::Script(ScriptCall {
                call_type,
                source_info: source_info.to_owned(),
            }),
            args: args.into_iter().collect(),
            source_info: source_info.to_owned(),
            current_line_offset: 0,
            current: None, // TODO(source-info): fill this out
            entry: None,   // TODO(source-info): fill this out
        });

        if matches!(call_type, ScriptCallType::Source) {
            self.script_source_depth += 1;
        }
    }

    /// Pushes a new trap handler frame onto the stack.
    ///
    /// # Arguments
    ///
    /// * `signal` - The signal being handled.
    /// * `handler` - The trap handler being invoked, if any.
    pub fn push_trap_handler(
        &mut self,
        signal: traps::TrapSignal,
        handler: Option<&traps::TrapHandler>,
    ) {
        let source_info =
            handler.map_or_else(crate::SourceInfo::default, |h| h.source_info.clone());

        self.frames.push_front(Frame {
            frame_type: FrameType::TrapHandler(signal),
            args: vec![],
            source_info,
            current_line_offset: 0,
            current: None, // TODO(source-info): fill this out
            entry: None,   // TODO(source-info): fill this out
        });

        self.active_trap_signals.insert(signal);
    }

    /// Pushes a new eval frame onto the stack.
    pub fn push_eval(&mut self) {
        self.frames.push_front(Frame {
            frame_type: FrameType::Eval,
            args: vec![],
            source_info: crate::SourceInfo::from("eval"), // TODO(source-info): fill this out
            current_line_offset: 0,
            current: None, // TODO(source-info): fill this out
            entry: None,   // TODO(source-info): fill this out
        });
    }

    /// Pushes a new command string frame onto the stack.
    pub fn push_command_string(&mut self) {
        self.frames.push_front(Frame {
            frame_type: FrameType::CommandString,
            args: vec![],
            source_info: crate::SourceInfo::from("environment"),
            current_line_offset: 0,
            current: None, // TODO(source-info): fill this out
            entry: None,   // TODO(source-info): fill this out
        });
    }

    /// Pushes a new interactive session frame onto the stack.
    pub fn push_interactive_session(&mut self) {
        self.frames.push_front(Frame {
            frame_type: FrameType::InteractiveSession,
            args: vec![],
            current_line_offset: 0,
            source_info: crate::SourceInfo::from("main"),
            current: None, // TODO(source-info): fill this out
            entry: None,   // TODO(source-info): fill this out
        });
    }

    /// Pushes a new function call frame onto the stack.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the function being called.
    /// * `function` - The function being called.
    /// * `args` - The positional arguments for the function call.
    pub fn push_function(
        &mut self,
        name: impl Into<String>,
        function: &functions::Registration,
        args: impl IntoIterator<Item = String>,
    ) {
        self.frames.push_front(Frame {
            frame_type: FrameType::Function(FunctionCall {
                function_name: name.into(),
                function: function.to_owned(),
            }),
            args: args.into_iter().collect(),
            source_info: function.source().clone(),
            entry: function.definition().location().map(|span| span.start),
            current: None, // TODO(source-info): fill this out
            current_line_offset: 0,
        });

        self.func_call_depth += 1;
    }

    /// Iterates through the function calls on the stack.
    pub fn iter_function_calls(&self) -> impl Iterator<Item = &FunctionCall> {
        self.iter().filter_map(|frame| {
            if let FrameType::Function(call) = &frame.frame_type {
                Some(call)
            } else {
                None
            }
        })
    }

    /// Iterates through the script calls on the stack.
    pub fn iter_script_calls(&self) -> impl Iterator<Item = &ScriptCall> {
        self.iter().filter_map(|frame| {
            if let FrameType::Script(call) = &frame.frame_type {
                Some(call)
            } else {
                None
            }
        })
    }

    /// Returns whether or not the current script stack frame is a sourced
    /// script.
    pub fn in_sourced_script(&self) -> bool {
        self.iter_script_calls()
            .next()
            .is_some_and(|call| matches!(call.call_type, ScriptCallType::Source))
    }

    /// Returns the current depth of function calls in the call stack.
    pub const fn function_call_depth(&self) -> usize {
        self.func_call_depth
    }

    /// Returns the current depth of sourced script calls in the call stack.
    pub const fn script_source_depth(&self) -> usize {
        self.script_source_depth
    }

    /// Returns whether the given trap signal is currently being handled
    /// (i.e., there is an active frame on the stack for this signal).
    pub fn is_trap_signal_active(&self, signal: traps::TrapSignal) -> bool {
        self.active_trap_signals.contains(&signal)
    }

    /// Clears the set of active trap signals. This should be called when
    /// creating subshells so they start with fresh trap execution state
    /// independent of the parent shell's currently-executing traps.
    pub fn clear_active_trap_signals(&mut self) {
        self.active_trap_signals.clear();
    }

    /// Returns whether the given trap signal is currently suppressed.
    pub const fn is_trap_delivery_suppressed(&self) -> bool {
        self.trap_delivery_suppress_count > 0
    }

    /// Acquires a block on trap delivery, preventing traps from being delivered
    /// until the block is released. Multiple blocks may be acquired, and
    /// trap delivery will remain suppressed until all blocks have been
    /// released.
    pub const fn acquire_trap_delivery_block(&mut self) {
        self.trap_delivery_suppress_count += 1;
    }

    /// Releases a block on trap delivery; note that trap delivery will remain
    /// suppressed until all blocks have been released.
    pub const fn release_trap_delivery_block(&mut self) {
        self.trap_delivery_suppress_count = self.trap_delivery_suppress_count.saturating_sub(1);
    }

    /// Returns whether or not the shell is actively executing in a shell
    /// function.
    pub fn in_function(&self) -> bool {
        self.iter_function_calls().next().is_some()
    }

    /// Returns the current depth of the call stack.
    pub fn depth(&self) -> usize {
        self.frames.len()
    }

    /// Returns whether or not the call stack is empty.
    pub fn is_empty(&self) -> bool {
        self.frames.is_empty()
    }

    /// Returns an iterator over the call frames, starting from the most
    /// recent.
    pub fn iter(&self) -> impl Iterator<Item = &Frame> {
        self.frames.iter()
    }

    /// Returns a mutable iterator over the call frames, starting from the most
    /// recent.
    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut Frame> {
        self.frames.iter_mut()
    }
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;

    use pretty_assertions::assert_matches;

    use super::*;
    use crate::SourceInfo;

    #[test]
    fn test_call_stack_new() {
        let stack = CallStack::new();
        assert!(stack.is_empty());
        assert_eq!(stack.depth(), 0);
    }

    #[test]
    fn test_call_stack_default() {
        let stack = CallStack::default();
        assert!(stack.is_empty());
        assert_eq!(stack.depth(), 0);
    }

    #[test]
    fn test_call_stack_push_pop() {
        let mut stack = CallStack::new();

        stack.push_script(
            ScriptCallType::Source,
            &SourceInfo::from(PathBuf::from("script1.sh")),
            vec![],
        );
        assert!(!stack.is_empty());
        assert_eq!(stack.depth(), 1);

        stack.push_script(
            ScriptCallType::Run,
            &SourceInfo::from(PathBuf::from("script2.sh")),
            vec![],
        );
        assert_eq!(stack.depth(), 2);

        let frame = stack.pop().unwrap();
        assert_matches!(
            frame.frame_type,
            FrameType::Script(ScriptCall {
                call_type: ScriptCallType::Run,
                source_info: SourceInfo {
                    source: file_path,
                    ..
                },
            }) if &file_path == "script2.sh"
        );
        assert_eq!(stack.depth(), 1);

        let frame = stack.pop().unwrap();
        assert_matches!(
            frame.frame_type,
            FrameType::Script(ScriptCall {
                call_type: ScriptCallType::Source,
                source_info: SourceInfo {
                    source: file_path,
                    ..
                },
            }) if &file_path == "script1.sh"
        );
        assert_eq!(stack.depth(), 0);
        assert!(stack.is_empty());
    }

    #[test]
    fn test_call_stack_pop_empty() {
        let mut stack = CallStack::new();
        assert!(stack.pop().is_none());
    }

    #[test]
    fn test_in_sourced_script() {
        let mut stack = CallStack::new();
        assert!(!stack.in_sourced_script());

        stack.push_script(
            ScriptCallType::Run,
            &SourceInfo::from(PathBuf::from("script1.sh")),
            vec![],
        );
        assert!(!stack.in_sourced_script());

        stack.push_script(
            ScriptCallType::Source,
            &SourceInfo::from(PathBuf::from("script2.sh")),
            vec![],
        );
        assert!(stack.in_sourced_script());

        stack.pop();
        assert!(!stack.in_sourced_script());
    }

    #[test]
    fn test_call_stack_iter() {
        let mut stack = CallStack::new();
        stack.push_script(
            ScriptCallType::Source,
            &SourceInfo::from(PathBuf::from("script1.sh")),
            vec![],
        );
        stack.push_script(
            ScriptCallType::Run,
            &SourceInfo::from(PathBuf::from("script2.sh")),
            vec![],
        );
        stack.push_script(
            ScriptCallType::Source,
            &SourceInfo::from(PathBuf::from("script3.sh")),
            vec![],
        );

        let frames: Vec<_> = stack.iter().collect();
        assert_eq!(frames.len(), 3);
        assert_matches!(&frames[0].frame_type, FrameType::Script(ScriptCall { source_info: SourceInfo { source: file_path, .. }, .. }) if file_path == "script3.sh");
        assert_matches!(&frames[1].frame_type, FrameType::Script(ScriptCall { source_info: SourceInfo { source: file_path, .. }, .. }) if file_path == "script2.sh");
        assert_matches!(&frames[2].frame_type, FrameType::Script(ScriptCall { source_info: SourceInfo { source: file_path, .. }, .. }) if file_path == "script1.sh");
    }
}