texlang/

variable.rs

//! Texlang variables API
//!
//! The documentation website contains
//!   [a tutorial specifically about Texlang's variables API](https://texcraft.dev/texlang/07-variables.html).
//! See that documentation for an overview of the API.
//! As usual, the documentation here is meant as reference.

use crate::command;
use crate::error;
use crate::parse::OptionalEquals;
use crate::prelude as txl;
use crate::token;
use crate::traits::*;
use crate::types;
use crate::vm;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fmt::Debug;
use std::hash::{Hash, Hasher};
use texcraft_stdext::collections::groupingmap;

/// Function signature for a variable's immutable getter.
///
/// In Texcraft all [Variable]s are built from an immutable getter and a mutable getter.
/// This type alias just defines the signature of immutable getters.
/// The first argument is the state, and the second argument is the index of the variable.
pub type RefFn<S, T> = fn(state: &S, index: Index) -> &T;

/// Function signature for a variable's mutable getters.
///
/// In Texcraft all [Variable]s are built from an immutable getter and a mutable getter.
/// This type alias just defines the signature of mutable getters.
/// The first argument is the state, and the second argument is the index of the variable.
pub type MutRefFn<S, T> = fn(state: &mut S, index: Index) -> &mut T;

/// Index of a variable within an array.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Default)]
pub struct Index(pub usize);

impl From<usize> for Index {
    fn from(value: usize) -> Self {
        Index(value)
    }
}

/// Specification for how the index of an array variable is determined.
///
/// Obtaining a variable from a command involves determining the variable's [Index].
/// This index is ultimately passed into the variable's getters to get a reference or mutable reference
/// to the underlying Rust value.
pub enum IndexResolver<S> {
    /// A static index, provided in the enum variant.
    ///
    /// This resolver is used for commands that point to a specific entry in a array.
    /// For example, after executing `\countdef\A 30`, the `\A` control sequence points
    /// at the count register with index 30.
    /// The command backing `\A` uses a static resolver with index 30.
    Static(Index),
    /// A dynamic index that is determined by reading the input token stream.
    ///
    /// For example, in `\count 4` the index of `4` is determined by parsing a number
    /// from the input token stream.
    Dynamic(fn(token::Token, &mut vm::ExpandedStream<S>) -> txl::Result<Index>),
}

impl<S> IndexResolver<S> {
    /// Determine the index of a variable using the input token stream.
    fn resolve(
        &self,
        token: token::Token,
        input: &mut vm::ExpandedStream<S>,
    ) -> txl::Result<Index> {
        match self {
            IndexResolver::Static(addr) => Ok(*addr),
            // TODO: we should add support for propagation information here.
            // We need to add to the execution stack
            IndexResolver::Dynamic(f) => f(token, input),
        }
    }
}

/// A TeX variable command.
///
/// Variable commands are _resolved_ to obtain a [Variable].
///
/// A command consists of three parts.
///
/// The first two parts are an immutable getter (of type [RefFn])
/// and a mutable getter (of type [MutRefFn]).
/// Both of these getters accept a reference to the state (where the variable's value lives)
/// and an [Index].
/// The index can be used by the getters to return different values in memory.
/// For example, if the getters read from an array, the index may be just the index in the array.
/// This mechanism allows users of Texlang to write one pair of getters that can be used in many variables.
///
/// The third part of a command is an [IndexResolver] that is used to determine the aforementioned index
/// of a variable at runtime.
/// The process of resolving a command involves determining the [Index] and returning a [Variable] type,
/// which under the hood is just the two getters and this index.
pub struct Command<S> {
    getters: Getters<S>,
    index_resolver: Option<IndexResolver<S>>,
}

impl<S> Command<S> {
    /// Create a new variable command.
    pub fn new_singleton<T: SupportedType>(
        ref_fn: RefFn<S, T>,
        ref_mut_fn: MutRefFn<S, T>,
    ) -> Command<S> {
        SupportedType::new_command(ref_fn, ref_mut_fn, None)
    }

    /// Create a new variable command.
    pub fn new_array<T: SupportedType>(
        ref_fn: RefFn<S, T>,
        ref_mut_fn: MutRefFn<S, T>,
        index_resolver: IndexResolver<S>,
    ) -> Command<S> {
        SupportedType::new_command(ref_fn, ref_mut_fn, Some(index_resolver))
    }

    /// Create a new getter provider.
    ///
    /// A getter provider is a variable command that is not intended to be invoked directly -
    ///     in fact, the variable command will panic the program if it is invoked.
    /// Instead the provider is included in a VM's initial commands so that
    ///     the VM has a reference to the getters inside the command.
    /// If a variable with the same getters is subsequently inserted into the commands map
    ///     (for example, by a `\countdef` type command), the VM can still serialize that
    ///     variable using the getters provided here.
    pub fn new_getter_provider<T: SupportedType>(
        ref_fn: RefFn<S, T>,
        ref_mut_fn: MutRefFn<S, T>,
    ) -> Command<S> {
        SupportedType::new_command(
            ref_fn,
            ref_mut_fn,
            Some(IndexResolver::Dynamic(|_, _| panic!())),
        )
    }

    /// Create a command that points to a specific element in the array referenced by this command.
    pub(crate) fn new_array_element(&self, index: Index) -> Self {
        Self {
            getters: self.getters.clone(),
            index_resolver: Some(IndexResolver::Static(index)),
        }
    }

    /// Returns true if the variable command corresponds to arithmetic variable(s) (e.g. integer or glue).
    ///
    /// TODO: we should just expose the type.
    pub fn is_arithmetic(&self) -> bool {
        match self.getters {
            Getters::Int(_, _)
            | Getters::SmallInt(_, _)
            | Getters::Dimen(_, _)
            | Getters::Glue(_, _) => true,
            Getters::CatCode(_, _)
            | Getters::MathCode(_, _)
            | Getters::TokenList(_, _)
            | Getters::Font(_, _) => false,
        }
    }
}

impl<S: TexlangState> Command<S> {
    /// Resolve the command to obtain a [Variable].
    pub fn resolve(
        &self,
        token: token::Token,
        input: &mut vm::ExpandedStream<S>,
    ) -> txl::Result<Variable<S>> {
        let index = match &self.index_resolver {
            None => Index(0),
            Some(index_resolver) => {
                input
                    .vm_mut()
                    .stack_push(token, error::OperationKind::VariableIndex);
                let err_or = index_resolver.resolve(token, input);
                input.vm_mut().stack_pop();
                err_or?
            }
        };
        Ok(new_variable(&self.getters, index))
    }

    /// Resolve the command to obtain a [Variable] of a specific type.
    pub fn resolve_type<T: SupportedType>(
        &self,
        token: token::Token,
        input: &mut vm::ExpandedStream<S>,
    ) -> txl::Result<Option<TypedVariable<S, T>>> {
        let Some((ref_fn, ref_mut_fn)) = T::try_cast(self) else {
            return Ok(None);
        };
        let index = match &self.index_resolver {
            None => Index(0),
            Some(index_resolver) => index_resolver.resolve(token, input)?,
        };
        Ok(Some(TypedVariable(ref_fn, ref_mut_fn, index)))
    }
}

impl<S> Command<S> {
    pub(crate) fn key(&self) -> CommandKey {
        let getters_key = self.getters.key();
        match &self.index_resolver {
            None => CommandKey::Singleton(getters_key),
            Some(index_resolver) => match index_resolver {
                IndexResolver::Static(a) => CommandKey::ArrayStatic(getters_key, *a),
                IndexResolver::Dynamic(f) => CommandKey::ArrayDynamic(getters_key, *f as usize),
            },
        }
    }
}

impl<S: TexlangState> Command<S> {
    /// Resolve the command to a variable and return the value of the variable.
    pub fn value<'a>(
        &self,
        token: token::Token,
        input: &'a mut vm::ExpandedStream<S>,
    ) -> txl::Result<ValueRef<'a>> {
        Ok(self.resolve(token, input)?.value(input.state()))
    }

    /// Resolve the command to a variable and set the value of the variable using the following tokens in the input stream.
    ///
    /// This function is used in TeX code like `\variable = 3`.
    /// In this case `\variable` is a command which resolves to a variable without consuming any more input.
    /// The variable is populated using the input `= 3` that follows.
    pub(crate) fn set_value_using_input(
        &self,
        token: token::Token,
        input: &mut vm::ExecutionInput<S>,
        scope: groupingmap::Scope,
    ) -> txl::Result<()> {
        let variable = self.resolve(token, input.as_mut())?;
        input
            .vm_mut()
            .stack_push(token, error::OperationKind::VariableAssignment);
        let err_or = variable.set_value_using_input(input, scope);
        input.vm_mut().stack_pop();
        err_or
    }
}

/// A key that uniquely identifies commands. If two commands have the same command key, they are the same command.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(crate) enum CommandKey {
    Singleton(GettersKey),
    ArrayStatic(GettersKey, Index),
    ArrayDynamic(GettersKey, usize),
}

impl CommandKey {
    pub(crate) fn getter_key(&self) -> GettersKey {
        match self {
            CommandKey::Singleton(k) => *k,
            CommandKey::ArrayStatic(k, _) => *k,
            CommandKey::ArrayDynamic(k, _) => *k,
        }
    }
}

/// A key that uniquely identifies the getters ([RefFn] and [MutRefFn]) in a command or variable.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(crate) struct GettersKey(usize, usize);

/// A TeX variable of a specific Rust type `T`.
pub struct TypedVariable<S, T>(RefFn<S, T>, MutRefFn<S, T>, Index);

impl<S, T> Copy for TypedVariable<S, T> {}

impl<S, T> Clone for TypedVariable<S, T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<S, T> TypedVariable<S, T> {
    /// Returns an immutable reference to the variable's value.
    pub fn get<'a>(&self, state: &'a S) -> &'a T {
        (self.0)(state, self.2)
    }

    fn key(&self) -> (usize, usize, Index) {
        (self.0 as usize, self.1 as usize, self.2)
    }
}

impl<S, T> TypedVariable<S, T>
where
    S: TexlangState,
    T: SupportedType,
{
    /// Sets the value of the variable.
    ///
    /// The input and scope must be passed because of TeX's grouping semantics.
    /// When the current group ends, any variable assignments made in the group
    ///     are rolled back.
    /// Thus this function generally saves the current value of the variable in the VM so that it
    ///     can be restored later.
    /// This is why the full input must be provided, and not just the state.
    pub fn set(&self, input: &mut vm::ExecutionInput<S>, scope: groupingmap::Scope, value: T) {
        let r: &mut T = (self.1)(input.state_mut(), self.2);
        let overwritten_value = std::mem::replace(r, value);
        // We guard the function call behind this conditional to make the current function small
        // to make it beneficial to inline it.
        if !input.groups().is_empty() {
            SupportedType::update_save_stack(input, self, scope, overwritten_value);
        } else {
            SupportedType::recycle(input, overwritten_value);
        }
    }
}
impl<S, T> TypedVariable<S, T>
where
    S: TexlangState,
    T: SupportedType,
{
    fn set_using_input(
        &self,
        input: &mut vm::ExecutionInput<S>,
        scope: groupingmap::Scope,
    ) -> txl::Result<()> {
        let (_, value) = <(OptionalEquals, T)>::parse(input)?;
        self.set(input, scope, value);
        Ok(())
    }
}

impl<S, T> PartialEq for TypedVariable<S, T> {
    fn eq(&self, rhs: &TypedVariable<S, T>) -> bool {
        self.key() == rhs.key()
    }
}

impl<S, T> Eq for TypedVariable<S, T> {}

impl<S, T> Hash for TypedVariable<S, T> {
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher,
    {
        self.key().hash(state);
    }
}

/// Trait satisfied by all Rust types that can be used as TeX variables.
///
/// It exists to make the variables API more ergonomic.
/// For example, it is used to provide a uniform constructor [Command::new_array] for commands.
/// The trait cannot be implemented for new types.
pub trait SupportedType: Sized + Debug + Parsable {
    /// Create a new command of this type with the provided reference functions and index resolver.
    fn new_command<S>(
        ref_fn: RefFn<S, Self>,
        ref_mut_fn: MutRefFn<S, Self>,
        index_resolver: Option<IndexResolver<S>>,
    ) -> Command<S>;

    /// Update the VM's save stack after a variable assignment.
    fn update_save_stack<S>(
        input: &mut vm::ExecutionInput<S>,
        variable: &TypedVariable<S, Self>,
        scope: groupingmap::Scope,
        overwritten_value: Self,
    ) {
        (_, _, _, _) = (input, variable, scope, overwritten_value);
    }

    /// Recycle a value that's about to be dropped.
    ///
    /// This method is intended for token lists, in which case the vector is returned
    /// to the buffers pool.
    fn recycle<S>(input: &mut vm::ExecutionInput<S>, overwritten_value: Self) {
        (_, _) = (input, overwritten_value)
    }

    /// Create a new typed variable of this type from the getters in the command and the provided index.
    ///
    /// Return `None` if the command has a different type to `Self`.
    fn new_typed_variable<S>(command: &Command<S>, index: Index) -> Option<TypedVariable<S, Self>>;

    /// Try to cast a variable command to this type.
    fn try_cast<S>(command: &Command<S>) -> Option<(RefFn<S, Self>, MutRefFn<S, Self>)>;
}

/// This function is used to implement the [SupportedType::update_save_stack] method.
///
/// It's a bit janky to implement the logic as a free function, rather than, say, as
///     a default trait implementation.
/// The main problem is that the implementation calls into the save stack, which is an
///     internal VM data structure.
/// Any way that does this on the trait directly requires making some of the save stack
///     type public, because [SupportedType] is public.
fn update_save_stack<S, T: Clone + SupportedType, F>(
    input: &mut vm::ExecutionInput<S>,
    variable: &TypedVariable<S, T>,
    scope: groupingmap::Scope,
    overwritten_value: T,
    map_getter: F,
) where
    F: Fn(&mut SaveStackElement<S>) -> &mut SaveStackMap<S, T>,
{
    match scope {
        groupingmap::Scope::Global => {
            let n = input.groups().len();
            for _ in 0..n {
                let group = &mut input.groups()[0];
                if let Some(stale_value) = map_getter(group).remove(variable) {
                    SupportedType::recycle(input, stale_value);
                }
            }
        }
        groupingmap::Scope::Local => {
            if let Some((group, _)) = input.current_group_mut() {
                if let Some(stale_value) = map_getter(group).save(*variable, overwritten_value) {
                    SupportedType::recycle(input, stale_value);
                }
            }
        }
    }
}

macro_rules! supported_type_impl {
    ( $(
        {
            rust_type: $rust_type: path,
            enum_variant: $enum_variant: ident,
            $( save_stack_field: $save_stack_field: ident, )?
            $( recycle_fn: $recycle_fn: ident, )?
        },
    )+ ) => {

        /// Immutable reference to the value of a variable.
        pub enum ValueRef<'a> {
            $(
                $enum_variant(&'a $rust_type),
            )+
        }

        /// TeX variable of any type.
        ///
        /// A variable uniquely identifies a Rust value in the state, like an `i32`.
        /// Operations on this value (like reading or setting the value) can be done in two ways:
        ///
        /// 1. (Easy, less flexible) Use the methods directly on this type like [Variable::value]
        ///    to read the value.
        ///    These methods are really ergonomic.
        ///    The problem with the value method specifically is that the result
        ///    is a reference which keeps the borrow of the state alive.
        ///    Thus, while holding onto the result of the value, you can't do anything this the
        ///    input stream like reading an argument.
        ///    This is especially a problem when you need to perform a different action depending on the concrete type of the variable.
        ///
        /// 2. (Trickier, more flexible) Match on the type's enum variants to determine the
        ///    concrete type of the variable.
        ///    The [TypedVariable] value obtained in this way can be used to perform operations on the value.
        ///    The main benefit of this approach is that after matching on the type, you can still use the input
        ///    stream to do things because there is not borrow alive.
        ///
        pub enum Variable<S> {
            $(
                $enum_variant(TypedVariable<S, $rust_type>),
            )+
        }

        fn new_variable<S>(getters: &Getters<S>, index: Index) -> Variable<S> {
            match getters {
                $(
                    Getters::$enum_variant(a, b) => Variable::$enum_variant(TypedVariable(*a, *b, index)),
                )+
            }
        }

        impl<S: TexlangState> Variable<S> {
            /// Return a reference to the value of the variable.
            pub fn value<'a>(&self, state: &'a S) -> ValueRef<'a> {
                match self {
                    $(
                        Variable::$enum_variant(variable) => ValueRef::$enum_variant(variable.get(state)),
                    )+
                }
            }

            /// Set the value of a variable using the following tokens in the input stream.
            fn set_value_using_input(
                &self,
                input: &mut vm::ExecutionInput<S>,
                scope: groupingmap::Scope,
            ) -> txl::Result<()> {
                match self {
                    $(
                        Variable::$enum_variant(variable) => variable.set_using_input(input, scope),
                    )+
                }
            }
        }

        enum Getters<S> {
            $(
                $enum_variant(RefFn<S, $rust_type>, MutRefFn<S, $rust_type>),
            )+
        }

        impl<S> Clone for Getters<S> {
            fn clone(&self) -> Self {
                match self {
                    $(
                        Self::$enum_variant(a, b) => Self::$enum_variant(*a, *b),
                    )+
                }
            }
        }

        impl<S> Getters<S> {
            fn key(&self) -> GettersKey {
                match self {
                    $(
                        Getters::$enum_variant(a, b) => GettersKey(*a as usize, *b as usize),
                    )+
                }
            }
        }

        $(
        impl SupportedType for $rust_type {
            fn new_command<S>(
                ref_fn: RefFn<S, Self>,
                ref_mut_fn: MutRefFn<S, Self>,
                index_resolver: Option<IndexResolver<S>>,
            ) -> Command<S> {
                Command {
                    getters: Getters::$enum_variant(ref_fn, ref_mut_fn),
                    index_resolver,
                }
            }
            $(
            fn update_save_stack<S>(
                input: &mut vm::ExecutionInput<S>,
                variable: &TypedVariable<S, Self>,
                scope: groupingmap::Scope,
                overwritten_value: Self,
            ) {
                update_save_stack(input, variable, scope, overwritten_value, |element| {
                    &mut element.$save_stack_field
                })
            }
            )?
            $(
            fn recycle<S>(input: &mut vm::ExecutionInput<S>, overwritten_value: Self) {
                $recycle_fn(input, overwritten_value)
            }
            )?
            fn new_typed_variable<S>(
                command: &Command<S>,
                index: Index,
            ) -> Option<TypedVariable<S, Self>> {
                match command.getters {
                    Getters::$enum_variant(a, b) => Some(TypedVariable(a, b, index)),
                    _ => None,
                }
            }
            fn try_cast<S>(command: &Command<S>) -> Option<(RefFn<S, Self>, MutRefFn<S, Self>)> {
                match command.getters {
                    Getters::$enum_variant(a, b) => Some((a,b)),
                    _ => None,
                }
            }
        }
        )+

        /// Internal VM data structure used to implement TeX's grouping semantics.
        pub(crate) struct SaveStackElement<S> {
            $( $(
                $save_stack_field: SaveStackMap<S, $rust_type>,
            )? )+
        }

        impl<S> Default for SaveStackElement<S> {
            fn default() -> Self {
                Self {
                    $( $(
                        $save_stack_field: Default::default(),
                    )? )+
                }
            }
        }

        impl<S> SaveStackElement<S> {
            pub(crate) fn restore(self, input: &mut vm::ExecutionInput<S>) {
                $( $(
                    self.$save_stack_field.restore(input);
                )? )+
            }

            pub(crate) fn serializable<'a>(
                &'a self,
                built_ins: &HashMap<GettersKey, token::CsName>,
            ) -> SerializableSaveStackElement<'a> {
                SerializableSaveStackElement {
                    $( $(
                        $save_stack_field: self.$save_stack_field.serializable(built_ins),
                    )? )+
                }
            }
        }

        #[cfg_attr(feature = "serde", derive(::serde::Serialize, ::serde::Deserialize))]
        pub(crate) struct SerializableSaveStackElement<'a> {
            $( $(
                $save_stack_field: Vec<(token::CsName, usize, Cow<'a, $rust_type>)>,
            )? )+
        }

        impl<'a> SerializableSaveStackElement<'a> {
            pub(crate) fn finish_deserialization<S>(
                self,
                built_ins: &HashMap<token::CsName, command::BuiltIn<S>>,
            ) -> SaveStackElement<S> {
                SaveStackElement {
                    $( $(
                        $save_stack_field: SaveStackMap::from_deserialized(self.$save_stack_field, built_ins),
                    )? )+
                }
            }
        }
    };
}

supported_type_impl!(
    {
        rust_type: i32,
        enum_variant: Int,
        save_stack_field: i32,
    },
    {
        rust_type: u8,
        enum_variant: SmallInt,
        save_stack_field: u8,
    },
    {
        rust_type: common::Scaled,
        enum_variant: Dimen,
        save_stack_field: dimen,
    },
    {
        rust_type: common::Glue,
        enum_variant: Glue,
        save_stack_field: glue,
    },
    {
        rust_type: types::CatCode,
        enum_variant: CatCode,
        save_stack_field: catcode,
    },
    {
        rust_type: types::MathCode,
        enum_variant: MathCode,
        save_stack_field: math_code,
    },
    {
        rust_type: Vec<token::Token>,
        enum_variant: TokenList,
        save_stack_field: token_list,
        recycle_fn: recycle_token_list,
    },
    {
        rust_type: types::Font,
        enum_variant: Font,
        save_stack_field: font,
    },
    // {
    //    rust_type: FontParam,
    //    enum_variant: FontParam,
    //    // no save_stack_field because font params are global
    // }
);

fn recycle_token_list<S>(input: &mut vm::ExecutionInput<S>, overwritten_value: Vec<token::Token>) {
    input.return_token_buffer(overwritten_value);
}

/// Internal VM data structure used to implement TeX's grouping semantics.
struct SaveStackMap<S, T>(HashMap<TypedVariable<S, T>, T>);

impl<S, T> Default for SaveStackMap<S, T> {
    fn default() -> Self {
        Self(HashMap::new())
    }
}

impl<S, T: Clone + SupportedType> SaveStackMap<S, T> {
    fn save(&mut self, variable: TypedVariable<S, T>, value: T) -> Option<T> {
        match self.0.entry(variable) {
            std::collections::hash_map::Entry::Occupied(_) => Some(value),
            std::collections::hash_map::Entry::Vacant(v) => {
                v.insert(value);
                None
            }
        }
    }

    fn remove(&mut self, variable: &TypedVariable<S, T>) -> Option<T> {
        self.0.remove(variable)
    }

    fn restore(self, input: &mut vm::ExecutionInput<S>) {
        for (v, restored_value) in self.0 {
            let dest = (v.1)(input.state_mut(), v.2);
            let overwritten_value = std::mem::replace(dest, restored_value);
            SupportedType::recycle(input, overwritten_value);
        }
    }

    fn serializable<'a>(
        &'a self,
        built_ins: &HashMap<GettersKey, token::CsName>,
    ) -> Vec<(token::CsName, usize, Cow<'a, T>)> {
        self.0
            .iter()
            .map(|(typed_variable, value): (&TypedVariable<S, T>, &T)| {
                let key = GettersKey(typed_variable.0 as usize, typed_variable.1 as usize);
                let cs_name = built_ins.get(&key).unwrap();
                (*cs_name, typed_variable.2 .0, Cow::Borrowed(value))
            })
            .collect()
    }
}

impl<S, T: SupportedType + Clone> SaveStackMap<S, T> {
    fn from_deserialized<'a>(
        deserialized: Vec<(token::CsName, usize, Cow<'a, T>)>,
        built_ins: &HashMap<token::CsName, command::BuiltIn<S>>,
    ) -> Self {
        let m = deserialized
            .into_iter()
            .map(
                |(cs_name, index, value): (token::CsName, usize, Cow<'a, T>)| {
                    // TODO: surface error here
                    let built_in = built_ins.get(&cs_name).unwrap();
                    let typed_variable = match built_in.cmd() {
                        command::Command::Variable(variable_command) => {
                            // TODO: error instead of unwrap
                            SupportedType::new_typed_variable(variable_command, Index(index))
                                .unwrap()
                        }
                        _ => panic!("wrong type of built in TODO return an error here"),
                    };
                    (typed_variable, value.into_owned())
                },
            )
            .collect();
        Self(m)
    }
}