Refactor `Scope` (#5797)

7 files changed, 222 insertions, 202 deletions

diff --git a/crates/typst-library/src/foundations/dict.rs b/crates/typst-library/src/foundations/dict.rs
index e4ab54e7..c93670c1 100644
--- a/crates/typst-library/src/foundations/dict.rs
+++ b/crates/typst-library/src/foundations/dict.rs
@@ -261,7 +261,12 @@ pub struct ToDict(Dict);
 
 cast! {
     ToDict,
-    v: Module => Self(v.scope().iter().map(|(k, v, _)| (Str::from(k.clone()), v.clone())).collect()),
+    v: Module => Self(v
+        .scope()
+        .iter()
+        .map(|(k, b)| (Str::from(k.clone()), b.read().clone()))
+        .collect()
+    ),
 }
 
 impl Debug for Dict {
diff --git a/crates/typst-library/src/foundations/func.rs b/crates/typst-library/src/foundations/func.rs
index a05deb1f..741b6633 100644
--- a/crates/typst-library/src/foundations/func.rs
+++ b/crates/typst-library/src/foundations/func.rs
@@ -259,7 +259,7 @@ impl Func {
         let scope =
             self.scope().ok_or("cannot access fields on user-defined functions")?;
         match scope.get(field) {
-            Some(field) => Ok(field),
+            Some(binding) => Ok(binding.read()),
             None => match self.name() {
                 Some(name) => bail!("function `{name}` does not contain field `{field}`"),
                 None => bail!("function does not contain field `{field}`"),
diff --git a/crates/typst-library/src/foundations/mod.rs b/crates/typst-library/src/foundations/mod.rs
index a790da4f..c335484f 100644
--- a/crates/typst-library/src/foundations/mod.rs
+++ b/crates/typst-library/src/foundations/mod.rs
@@ -94,7 +94,7 @@ pub static FOUNDATIONS: Category;
 
 /// Hook up all `foundations` definitions.
 pub(super) fn define(global: &mut Scope, inputs: Dict, features: &Features) {
-    global.category(FOUNDATIONS);
+    global.start_category(FOUNDATIONS);
     global.define_type::<bool>();
     global.define_type::<i64>();
     global.define_type::<f64>();
@@ -301,7 +301,7 @@ pub fn eval(
     let dict = scope;
     let mut scope = Scope::new();
     for (key, value) in dict {
-        scope.define_spanned(key, value, span);
+        scope.bind(key.into(), Binding::new(value, span));
     }
     (engine.routines.eval_string)(engine.routines, engine.world, &text, span, mode, scope)
 }
diff --git a/crates/typst-library/src/foundations/module.rs b/crates/typst-library/src/foundations/module.rs
index 3ee59c10..3259c17e 100644
--- a/crates/typst-library/src/foundations/module.rs
+++ b/crates/typst-library/src/foundations/module.rs
@@ -4,7 +4,7 @@ use std::sync::Arc;
 use ecow::{eco_format, EcoString};
 use typst_syntax::FileId;
 
-use crate::diag::StrResult;
+use crate::diag::{bail, StrResult};
 use crate::foundations::{repr, ty, Content, Scope, Value};
 
 /// An module of definitions.
@@ -118,11 +118,14 @@ impl Module {
     }
 
     /// Try to access a definition in the module.
-    pub fn field(&self, name: &str) -> StrResult<&Value> {
-        self.scope().get(name).ok_or_else(|| match &self.name {
-            Some(module) => eco_format!("module `{module}` does not contain `{name}`"),
-            None => eco_format!("module does not contain `{name}`"),
-        })
+    pub fn field(&self, field: &str) -> StrResult<&Value> {
+        match self.scope().get(field) {
+            Some(binding) => Ok(binding.read()),
+            None => match &self.name {
+                Some(name) => bail!("module `{name}` does not contain `{field}`"),
+                None => bail!("module does not contain `{field}`"),
+            },
+        }
     }
 
     /// Extract the module's content.
diff --git a/crates/typst-library/src/foundations/plugin.rs b/crates/typst-library/src/foundations/plugin.rs
index cbc0f52d..a33f1cb9 100644
--- a/crates/typst-library/src/foundations/plugin.rs
+++ b/crates/typst-library/src/foundations/plugin.rs
@@ -8,7 +8,7 @@ use wasmi::Memory;
 
 use crate::diag::{bail, At, SourceResult, StrResult};
 use crate::engine::Engine;
-use crate::foundations::{cast, func, scope, Bytes, Func, Module, Scope, Value};
+use crate::foundations::{cast, func, scope, Binding, Bytes, Func, Module, Scope, Value};
 use crate::loading::{DataSource, Load};
 
 /// Loads a WebAssembly module.
@@ -369,7 +369,7 @@ impl Plugin {
             if matches!(export.ty(), wasmi::ExternType::Func(_)) {
                 let name = EcoString::from(export.name());
                 let func = PluginFunc { plugin: shared.clone(), name: name.clone() };
-                scope.define(name, Func::from(func));
+                scope.bind(name, Binding::detached(Func::from(func)));
             }
         }
 
diff --git a/crates/typst-library/src/foundations/scope.rs b/crates/typst-library/src/foundations/scope.rs
index b7b4a6d9..e73afeac 100644
--- a/crates/typst-library/src/foundations/scope.rs
+++ b/crates/typst-library/src/foundations/scope.rs
@@ -5,8 +5,8 @@ use std::fmt::{self, Debug, Formatter};
 use std::hash::{Hash, Hasher};
 
 use ecow::{eco_format, EcoString};
+use indexmap::map::Entry;
 use indexmap::IndexMap;
-use typst_syntax::ast::{self, AstNode};
 use typst_syntax::Span;
 use typst_utils::Static;
 
@@ -46,14 +46,14 @@ impl<'a> Scopes<'a> {
         self.top = self.scopes.pop().expect("no pushed scope");
     }
 
-    /// Try to access a variable immutably.
-    pub fn get(&self, var: &str) -> HintedStrResult<&Value> {
+    /// Try to access a binding immutably.
+    pub fn get(&self, var: &str) -> HintedStrResult<&Binding> {
         std::iter::once(&self.top)
             .chain(self.scopes.iter().rev())
             .find_map(|scope| scope.get(var))
             .or_else(|| {
                 self.base.and_then(|base| match base.global.scope().get(var) {
-                    Some(value) => Some(value),
+                    Some(binding) => Some(binding),
                     None if var == "std" => Some(&base.std),
                     None => None,
                 })
@@ -61,14 +61,28 @@ impl<'a> Scopes<'a> {
             .ok_or_else(|| unknown_variable(var))
     }
 
-    /// Try to access a variable immutably in math.
-    pub fn get_in_math(&self, var: &str) -> HintedStrResult<&Value> {
+    /// Try to access a binding mutably.
+    pub fn get_mut(&mut self, var: &str) -> HintedStrResult<&mut Binding> {
+        std::iter::once(&mut self.top)
+            .chain(&mut self.scopes.iter_mut().rev())
+            .find_map(|scope| scope.get_mut(var))
+            .ok_or_else(|| {
+                match self.base.and_then(|base| base.global.scope().get(var)) {
+                    Some(_) => cannot_mutate_constant(var),
+                    _ if var == "std" => cannot_mutate_constant(var),
+                    _ => unknown_variable(var),
+                }
+            })
+    }
+
+    /// Try to access a binding immutably in math.
+    pub fn get_in_math(&self, var: &str) -> HintedStrResult<&Binding> {
         std::iter::once(&self.top)
             .chain(self.scopes.iter().rev())
             .find_map(|scope| scope.get(var))
             .or_else(|| {
                 self.base.and_then(|base| match base.math.scope().get(var) {
-                    Some(value) => Some(value),
+                    Some(binding) => Some(binding),
                     None if var == "std" => Some(&base.std),
                     None => None,
                 })
@@ -81,20 +95,6 @@ impl<'a> Scopes<'a> {
             })
     }
 
-    /// Try to access a variable mutably.
-    pub fn get_mut(&mut self, var: &str) -> HintedStrResult<&mut Value> {
-        std::iter::once(&mut self.top)
-            .chain(&mut self.scopes.iter_mut().rev())
-            .find_map(|scope| scope.get_mut(var))
-            .ok_or_else(|| {
-                match self.base.and_then(|base| base.global.scope().get(var)) {
-                    Some(_) => cannot_mutate_constant(var),
-                    _ if var == "std" => cannot_mutate_constant(var),
-                    _ => unknown_variable(var),
-                }
-            })?
-    }
-
     /// Check if an std variable is shadowed.
     pub fn check_std_shadowed(&self, var: &str) -> bool {
         self.base.is_some_and(|base| base.global.scope().get(var).is_some())
@@ -104,84 +104,28 @@ impl<'a> Scopes<'a> {
     }
 }
 
-#[cold]
-fn cannot_mutate_constant(var: &str) -> HintedString {
-    eco_format!("cannot mutate a constant: {}", var).into()
-}
-
-/// The error message when a variable is not found.
-#[cold]
-fn unknown_variable(var: &str) -> HintedString {
-    let mut res = HintedString::new(eco_format!("unknown variable: {}", var));
-
-    if var.contains('-') {
-        res.hint(eco_format!(
-            "if you meant to use subtraction, try adding spaces around the minus sign{}: `{}`",
-            if var.matches('-').count() > 1 { "s" } else { "" },
-            var.replace('-', " - ")
-        ));
-    }
-
-    res
-}
-
-#[cold]
-fn unknown_variable_math(var: &str, in_global: bool) -> HintedString {
-    let mut res = HintedString::new(eco_format!("unknown variable: {}", var));
-
-    if matches!(var, "none" | "auto" | "false" | "true") {
-        res.hint(eco_format!(
-            "if you meant to use a literal, try adding a hash before it: `#{var}`",
-        ));
-    } else if in_global {
-        res.hint(eco_format!(
-            "`{var}` is not available directly in math, try adding a hash before it: `#{var}`",
-        ));
-    } else {
-        res.hint(eco_format!(
-            "if you meant to display multiple letters as is, try adding spaces between each letter: `{}`",
-            var.chars()
-                .flat_map(|c| [' ', c])
-                .skip(1)
-                .collect::<EcoString>()
-        ));
-        res.hint(eco_format!(
-            "or if you meant to display this as text, try placing it in quotes: `\"{var}\"`"
-        ));
-    }
-
-    res
-}
-
 /// A map from binding names to values.
 #[derive(Default, Clone)]
 pub struct Scope {
-    map: IndexMap<EcoString, Slot>,
+    map: IndexMap<EcoString, Binding>,
     deduplicate: bool,
     category: Option<Category>,
 }
 
+/// Scope construction.
 impl Scope {
     /// Create a new empty scope.
     pub fn new() -> Self {
         Default::default()
     }
 
-    /// Create a new scope with the given capacity.
-    pub fn with_capacity(capacity: usize) -> Self {
-        Self {
-            map: IndexMap::with_capacity(capacity),
-            ..Default::default()
-        }
-    }
-
     /// Create a new scope with duplication prevention.
     pub fn deduplicating() -> Self {
         Self { deduplicate: true, ..Default::default() }
     }
 
     /// Enter a new category.
-    pub fn category(&mut self, category: Category) {
+    pub fn start_category(&mut self, category: Category) {
         self.category = Some(category);
     }
 
@@ -190,102 +134,87 @@ impl Scope {
         self.category = None;
     }
 
-    /// Bind a value to a name.
-    #[track_caller]
-    pub fn define(&mut self, name: impl Into<EcoString>, value: impl IntoValue) {
-        self.define_spanned(name, value, Span::detached())
-    }
-
-    /// Bind a value to a name defined by an identifier.
-    #[track_caller]
-    pub fn define_ident(&mut self, ident: ast::Ident, value: impl IntoValue) {
-        self.define_spanned(ident.get().clone(), value, ident.span())
-    }
-
-    /// Bind a value to a name.
-    #[track_caller]
-    pub fn define_spanned(
-        &mut self,
-        name: impl Into<EcoString>,
-        value: impl IntoValue,
-        span: Span,
-    ) {
-        let name = name.into();
-
-        #[cfg(debug_assertions)]
-        if self.deduplicate && self.map.contains_key(&name) {
-            panic!("duplicate definition: {name}");
-        }
-
-        self.map.insert(
-            name,
-            Slot::new(value.into_value(), span, Kind::Normal, self.category),
-        );
-    }
-
-    /// Define a captured, immutable binding.
-    pub fn define_captured(
-        &mut self,
-        name: EcoString,
-        value: Value,
-        capturer: Capturer,
-        span: Span,
-    ) {
-        self.map.insert(
-            name,
-            Slot::new(value.into_value(), span, Kind::Captured(capturer), self.category),
-        );
-    }
-
     /// Define a native function through a Rust type that shadows the function.
-    pub fn define_func<T: NativeFunc>(&mut self) {
+    #[track_caller]
+    pub fn define_func<T: NativeFunc>(&mut self) -> &mut Binding {
         let data = T::data();
-        self.define(data.name, Func::from(data));
+        self.define(data.name, Func::from(data))
     }
 
     /// Define a native function with raw function data.
-    pub fn define_func_with_data(&mut self, data: &'static NativeFuncData) {
-        self.define(data.name, Func::from(data));
+    #[track_caller]
+    pub fn define_func_with_data(
+        &mut self,
+        data: &'static NativeFuncData,
+    ) -> &mut Binding {
+        self.define(data.name, Func::from(data))
     }
 
     /// Define a native type.
-    pub fn define_type<T: NativeType>(&mut self) {
+    #[track_caller]
+    pub fn define_type<T: NativeType>(&mut self) -> &mut Binding {
         let data = T::data();
-        self.define(data.name, Type::from(data));
+        self.define(data.name, Type::from(data))
     }
 
     /// Define a native element.
-    pub fn define_elem<T: NativeElement>(&mut self) {
+    #[track_caller]
+    pub fn define_elem<T: NativeElement>(&mut self) -> &mut Binding {
         let data = T::data();
-        self.define(data.name, Element::from(data));
+        self.define(data.name, Element::from(data))
     }
 
-    /// Try to access a variable immutably.
-    pub fn get(&self, var: &str) -> Option<&Value> {
-        self.map.get(var).map(Slot::read)
+    /// Define a built-in with compile-time known name and returns a mutable
+    /// reference to it.
+    ///
+    /// When the name isn't compile-time known, you should instead use:
+    /// - `Vm::bind` if you already have [`Binding`]
+    /// - `Vm::define`  if you only have a [`Value`]
+    /// - [`Scope::bind`](Self::bind) if you are not operating in the context of
+    ///   a `Vm` or if you are binding to something that is not an AST
+    ///   identifier (e.g. when constructing a dynamic
+    ///   [`Module`](super::Module))
+    #[track_caller]
+    pub fn define(&mut self, name: &'static str, value: impl IntoValue) -> &mut Binding {
+        #[cfg(debug_assertions)]
+        if self.deduplicate && self.map.contains_key(name) {
+            panic!("duplicate definition: {name}");
+        }
+
+        let mut binding = Binding::detached(value);
+        binding.category = self.category;
+        self.bind(name.into(), binding)
     }
+}
 
-    /// Try to access a variable mutably.
-    pub fn get_mut(&mut self, var: &str) -> Option<HintedStrResult<&mut Value>> {
-        self.map
-            .get_mut(var)
-            .map(Slot::write)
-            .map(|res| res.map_err(HintedString::from))
+/// Scope manipulation and access.
+impl Scope {
+    /// Inserts a binding into this scope and returns a mutable reference to it.
+    ///
+    /// Prefer `Vm::bind` if you are operating in the context of a `Vm`.
+    pub fn bind(&mut self, name: EcoString, binding: Binding) -> &mut Binding {
+        match self.map.entry(name) {
+            Entry::Occupied(mut entry) => {
+                entry.insert(binding);
+                entry.into_mut()
+            }
+            Entry::Vacant(entry) => entry.insert(binding),
+        }
     }
 
-    /// Get the span of a definition.
-    pub fn get_span(&self, var: &str) -> Option<Span> {
-        Some(self.map.get(var)?.span)
+    /// Try to access a binding immutably.
+    pub fn get(&self, var: &str) -> Option<&Binding> {
+        self.map.get(var)
     }
 
-    /// Get the category of a definition.
-    pub fn get_category(&self, var: &str) -> Option<Category> {
-        self.map.get(var)?.category
+    /// Try to access a binding mutably.
+    pub fn get_mut(&mut self, var: &str) -> Option<&mut Binding> {
+        self.map.get_mut(var)
     }
 
     /// Iterate over all definitions.
-    pub fn iter(&self) -> impl Iterator<Item = (&EcoString, &Value, Span)> {
-        self.map.iter().map(|(k, v)| (k, v.read(), v.span))
+    pub fn iter(&self) -> impl Iterator<Item = (&EcoString, &Binding)> {
+        self.map.iter()
     }
 }
 
@@ -318,64 +247,92 @@ pub trait NativeScope {
     fn scope() -> Scope;
 }
 
-/// A slot where a value is stored.
-#[derive(Clone, Hash)]
-struct Slot {
-    /// The stored value.
+/// A bound value with metadata.
+#[derive(Debug, Clone, Hash)]
+pub struct Binding {
+    /// The bound value.
     value: Value,
-    /// The kind of slot, determines how the value can be accessed.
-    kind: Kind,
-    /// A span associated with the stored value.
+    /// The kind of binding, determines how the value can be accessed.
+    kind: BindingKind,
+    /// A span associated with the binding.
     span: Span,
-    /// The category of the slot.
+    /// The category of the binding.
     category: Option<Category>,
 }
 
 /// The different kinds of slots.
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
-enum Kind {
+enum BindingKind {
     /// A normal, mutable binding.
     Normal,
     /// A captured copy of another variable.
     Captured(Capturer),
 }
 
-/// What the variable was captured by.
-#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
-pub enum Capturer {
-    /// Captured by a function / closure.
-    Function,
-    /// Captured by a context expression.
-    Context,
-}
+impl Binding {
+    /// Create a new binding with a span marking its definition site.
+    pub fn new(value: impl IntoValue, span: Span) -> Self {
+        Self {
+            value: value.into_value(),
+            span,
+            kind: BindingKind::Normal,
+            category: None,
+        }
+    }
 
-impl Slot {
-    /// Create a new slot.
-    fn new(value: Value, span: Span, kind: Kind, category: Option<Category>) -> Self {
-        Self { value, span, kind, category }
+    /// Create a binding without a span.
+    pub fn detached(value: impl IntoValue) -> Self {
+        Self::new(value, Span::detached())
     }
 
     /// Read the value.
-    fn read(&self) -> &Value {
+    pub fn read(&self) -> &Value {
         &self.value
     }
 
     /// Try to write to the value.
-    fn write(&mut self) -> StrResult<&mut Value> {
+    ///
+    /// This fails if the value is a read-only closure capture.
+    pub fn write(&mut self) -> StrResult<&mut Value> {
         match self.kind {
-            Kind::Normal => Ok(&mut self.value),
-            Kind::Captured(capturer) => {
-                bail!(
-                    "variables from outside the {} are \
-                     read-only and cannot be modified",
-                    match capturer {
-                        Capturer::Function => "function",
-                        Capturer::Context => "context expression",
-                    }
-                )
-            }
+            BindingKind::Normal => Ok(&mut self.value),
+            BindingKind::Captured(capturer) => bail!(
+                "variables from outside the {} are \
+                 read-only and cannot be modified",
+                match capturer {
+                    Capturer::Function => "function",
+                    Capturer::Context => "context expression",
+                }
+            ),
+        }
+    }
+
+    /// Create a copy of the binding for closure capturing.
+    pub fn capture(&self, capturer: Capturer) -> Self {
+        Self {
+            kind: BindingKind::Captured(capturer),
+            ..self.clone()
         }
     }
+
+    /// A span associated with the stored value.
+    pub fn span(&self) -> Span {
+        self.span
+    }
+
+    /// The category of the value, if any.
+    pub fn category(&self) -> Option<Category> {
+        self.category
+    }
+}
+
+/// What the variable was captured by.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
+pub enum Capturer {
+    /// Captured by a function / closure.
+    Function,
+    /// Captured by a context expression.
+    Context,
 }
 
 /// A group of related definitions.
@@ -417,3 +374,57 @@ pub struct CategoryData {
     pub title: &'static str,
     pub docs: &'static str,
 }
+
+/// The error message when trying to mutate a variable from the standard
+/// library.
+#[cold]
+fn cannot_mutate_constant(var: &str) -> HintedString {
+    eco_format!("cannot mutate a constant: {}", var).into()
+}
+
+/// The error message when a variable wasn't found.
+#[cold]
+fn unknown_variable(var: &str) -> HintedString {
+    let mut res = HintedString::new(eco_format!("unknown variable: {}", var));
+
+    if var.contains('-') {
+        res.hint(eco_format!(
+            "if you meant to use subtraction, \
+             try adding spaces around the minus sign{}: `{}`",
+            if var.matches('-').count() > 1 { "s" } else { "" },
+            var.replace('-', " - ")
+        ));
+    }
+
+    res
+}
+
+/// The error message when a variable wasn't found it math.
+#[cold]
+fn unknown_variable_math(var: &str, in_global: bool) -> HintedString {
+    let mut res = HintedString::new(eco_format!("unknown variable: {}", var));
+
+    if matches!(var, "none" | "auto" | "false" | "true") {
+        res.hint(eco_format!(
+            "if you meant to use a literal, \
+             try adding a hash before it: `#{var}`",
+        ));
+    } else if in_global {
+        res.hint(eco_format!(
+            "`{var}` is not available directly in math, \
+             try adding a hash before it: `#{var}`",
+        ));
+    } else {
+        res.hint(eco_format!(
+            "if you meant to display multiple letters as is, \
+             try adding spaces between each letter: `{}`",
+            var.chars().flat_map(|c| [' ', c]).skip(1).collect::<EcoString>()
+        ));
+        res.hint(eco_format!(
+            "or if you meant to display this as text, \
+             try placing it in quotes: `\"{var}\"`"
+        ));
+    }
+
+    res
+}
diff --git a/crates/typst-library/src/foundations/ty.rs b/crates/typst-library/src/foundations/ty.rs
index 973c1cb6..09f5efa1 100644
--- a/crates/typst-library/src/foundations/ty.rs
+++ b/crates/typst-library/src/foundations/ty.rs
@@ -8,7 +8,7 @@ use std::sync::LazyLock;
 use ecow::{eco_format, EcoString};
 use typst_utils::Static;
 
-use crate::diag::StrResult;
+use crate::diag::{bail, StrResult};
 use crate::foundations::{
     cast, func, AutoValue, Func, NativeFuncData, NoneValue, Repr, Scope, Value,
 };
@@ -95,9 +95,10 @@ impl Type {
 
     /// Get a field from this type's scope, if possible.
     pub fn field(&self, field: &str) -> StrResult<&'static Value> {
-        self.scope()
-            .get(field)
-            .ok_or_else(|| eco_format!("type {self} does not contain field `{field}`"))
+        match self.scope().get(field) {
+            Some(binding) => Ok(binding.read()),
+            None => bail!("type {self} does not contain field `{field}`"),
+        }
     }
 }