Accept closures for heading styling

6 files changed, 140 insertions, 35 deletions

diff --git a/src/diag.rs b/src/diag.rs
index 70a5bdaf..329e9a8a 100644
--- a/src/diag.rs
+++ b/src/diag.rs
@@ -108,3 +108,13 @@ impl<T> Trace<T> for TypResult<T> {
         })
     }
 }
+
+/// Transform `expected X, found Y` into `expected X or A, found Y`.
+pub fn with_alternative(msg: String, alt: &str) -> String {
+    let mut parts = msg.split(", found ");
+    if let (Some(a), Some(b)) = (parts.next(), parts.next()) {
+        format!("{} or {}, found {}", a, alt, b)
+    } else {
+        msg
+    }
+}
diff --git a/src/eval/value.rs b/src/eval/value.rs
index fffa729a..2de210d5 100644
--- a/src/eval/value.rs
+++ b/src/eval/value.rs
@@ -531,7 +531,7 @@ impl<T: Cast> Cast for Smart<T> {
 }
 
 /// Transform `expected X, found Y` into `expected X or A, found Y`.
-fn with_alternative(msg: String, alt: &str) -> String {
+pub fn with_alternative(msg: String, alt: &str) -> String {
     let mut parts = msg.split(", found ");
     if let (Some(a), Some(b)) = (parts.next(), parts.next()) {
         format!("{} or {}, found {}", a, alt, b)
diff --git a/src/library/heading.rs b/src/library/heading.rs
index 2044355e..dd78b147 100644
--- a/src/library/heading.rs
+++ b/src/library/heading.rs
@@ -15,24 +15,25 @@ pub struct HeadingNode {
 
 #[class]
 impl HeadingNode {
-    /// The heading's font family.
-    pub const FAMILY: Smart<FontFamily> = Smart::Auto;
-    /// The size of text in the heading. Just the surrounding text size if
-    /// `auto`.
-    pub const SIZE: Smart<Linear> = Smart::Auto;
-    /// The fill color of text in the heading. Just the surrounding text color
-    /// if `auto`.
-    pub const FILL: Smart<Paint> = Smart::Auto;
+    /// The heading's font family. Just the normal text family if `auto`.
+    pub const FAMILY: Leveled<Smart<FontFamily>> = Leveled::Value(Smart::Auto);
+    /// The color of text in the heading. Just the normal text color if `auto`.
+    pub const FILL: Leveled<Smart<Paint>> = Leveled::Value(Smart::Auto);
+    /// The size of text in the heading.
+    pub const SIZE: Leveled<Linear> = Leveled::Mapping(|level| {
+        let upscale = (1.6 - 0.1 * level as f64).max(0.75);
+        Relative::new(upscale).into()
+    });
     /// Whether text in the heading is strengthend.
-    pub const STRONG: bool = true;
+    pub const STRONG: Leveled<bool> = Leveled::Value(true);
     /// Whether text in the heading is emphasized.
-    pub const EMPH: bool = false;
+    pub const EMPH: Leveled<bool> = Leveled::Value(false);
     /// Whether the heading is underlined.
-    pub const UNDERLINE: bool = false;
+    pub const UNDERLINE: Leveled<bool> = Leveled::Value(false);
     /// The extra padding above the heading.
-    pub const ABOVE: Length = Length::zero();
+    pub const ABOVE: Leveled<Length> = Leveled::Value(Length::zero());
     /// The extra padding below the heading.
-    pub const BELOW: Length = Length::zero();
+    pub const BELOW: Leveled<Length> = Leveled::Value(Length::zero());
 
     fn construct(_: &mut Vm, args: &mut Args) -> TypResult<Template> {
         Ok(Template::show(Self {
@@ -43,52 +44,51 @@ impl HeadingNode {
 }
 
 impl Show for HeadingNode {
-    fn show(&self, _: &mut Vm, styles: StyleChain) -> TypResult<Template> {
-        let mut map = StyleMap::new();
+    fn show(&self, vm: &mut Vm, styles: StyleChain) -> TypResult<Template> {
+        macro_rules! resolve {
+            ($key:expr) => {
+                styles.get_cloned($key).resolve(vm, self.level)?
+            };
+        }
 
-        let upscale = (1.6 - 0.1 * self.level as f64).max(0.75);
-        map.set(
-            TextNode::SIZE,
-            styles.get(Self::SIZE).unwrap_or(Relative::new(upscale).into()),
-        );
+        let mut map = StyleMap::new();
+        map.set(TextNode::SIZE, resolve!(Self::SIZE));
 
-        if let Smart::Custom(family) = styles.get_ref(Self::FAMILY) {
+        if let Smart::Custom(family) = resolve!(Self::FAMILY) {
             map.set(
                 TextNode::FAMILY,
                 std::iter::once(family)
-                    .chain(styles.get_ref(TextNode::FAMILY))
-                    .cloned()
+                    .chain(styles.get_ref(TextNode::FAMILY).iter().cloned())
                     .collect(),
             );
         }
 
-        if let Smart::Custom(fill) = styles.get(Self::FILL) {
+        if let Smart::Custom(fill) = resolve!(Self::FILL) {
             map.set(TextNode::FILL, fill);
         }
 
-        if styles.get(Self::STRONG) {
+        if resolve!(Self::STRONG) {
             map.set(TextNode::STRONG, true);
         }
 
-        if styles.get(Self::EMPH) {
+        if resolve!(Self::EMPH) {
             map.set(TextNode::EMPH, true);
         }
 
+        let mut seq = vec![];
         let mut body = self.body.clone();
-        if styles.get(Self::UNDERLINE) {
+        if resolve!(Self::UNDERLINE) {
             body = body.underlined();
         }
 
-        let mut seq = vec![];
-
-        let above = styles.get(Self::ABOVE);
+        let above = resolve!(Self::ABOVE);
         if !above.is_zero() {
             seq.push(Template::Vertical(above.into()));
         }
 
         seq.push(body);
 
-        let below = styles.get(Self::BELOW);
+        let below = resolve!(Self::BELOW);
         if !below.is_zero() {
             seq.push(Template::Vertical(below.into()));
         }
@@ -98,3 +98,50 @@ impl Show for HeadingNode {
         ))
     }
 }
+
+/// Either the value or a closure mapping to the value.
+#[derive(Debug, Clone, PartialEq, Hash)]
+pub enum Leveled<T> {
+    /// A bare value.
+    Value(T),
+    /// A simple mapping from a heading level to a value.
+    Mapping(fn(usize) -> T),
+    /// A closure mapping from a heading level to a value.
+    Func(Func, Span),
+}
+
+impl<T: Cast> Leveled<T> {
+    /// Resolve the value based on the level.
+    pub fn resolve(self, vm: &mut Vm, level: usize) -> TypResult<T> {
+        match self {
+            Self::Value(value) => Ok(value),
+            Self::Mapping(mapping) => Ok(mapping(level)),
+            Self::Func(func, span) => {
+                let args = Args {
+                    span,
+                    items: vec![Arg {
+                        span,
+                        name: None,
+                        value: Spanned::new(Value::Int(level as i64), span),
+                    }],
+                };
+                func.call(vm, args)?.cast().at(span)
+            }
+        }
+    }
+}
+
+impl<T: Cast> Cast<Spanned<Value>> for Leveled<T> {
+    fn is(value: &Spanned<Value>) -> bool {
+        matches!(&value.v, Value::Func(_)) || T::is(&value.v)
+    }
+
+    fn cast(value: Spanned<Value>) -> StrResult<Self> {
+        match value.v {
+            Value::Func(f) => Ok(Self::Func(f, value.span)),
+            v => T::cast(v)
+                .map(Self::Value)
+                .map_err(|msg| with_alternative(msg, "function")),
+        }
+    }
+}
diff --git a/src/library/mod.rs b/src/library/mod.rs
index d90c96da..707264d9 100644
--- a/src/library/mod.rs
+++ b/src/library/mod.rs
@@ -62,10 +62,10 @@ pub mod prelude {
 
     pub use typst_macros::class;
 
-    pub use crate::diag::{At, TypResult};
+    pub use crate::diag::{with_alternative, At, StrResult, TypResult};
     pub use crate::eval::{
-        Args, Construct, Merge, Property, Scope, Set, Show, ShowNode, Smart, StyleChain,
-        StyleMap, StyleVec, Template, Value,
+        Arg, Args, Cast, Construct, Func, Merge, Property, Scope, Set, Show, ShowNode,
+        Smart, StyleChain, StyleMap, StyleVec, Template, Value,
     };
     pub use crate::frame::*;
     pub use crate::geom::*;
diff --git a/tests/ref/style/closure.png b/tests/ref/style/closure.png
new file mode 100644
index 00000000..77d88d0c
--- /dev/null
+++ b/tests/ref/style/closure.png
diff --git a/tests/typ/style/closure.typ b/tests/typ/style/closure.typ
new file mode 100644
index 00000000..2c9edd81
--- /dev/null
+++ b/tests/typ/style/closure.typ
@@ -0,0 +1,48 @@
+// Test styles with closure.
+
+---
+#set heading(
+  size: 10pt,
+  fill: lvl => if even(lvl) { red } else { blue },
+)
+
+= Heading 1
+== Heading 2
+=== Heading 3
+==== Heading 4
+
+---
+// Test in constructor.
+#heading(
+  level: 3,
+  size: 10pt,
+  strong: lvl => {
+    assert(lvl == 3)
+    false
+  }
+)[Level 3]
+
+---
+// Error: 22-26 expected font family or auto or function, found length
+#set heading(family: 10pt)
+= Heading
+
+---
+// Error: 29-38 cannot add integer and string
+#set heading(strong: lvl => lvl + "2")
+= Heading
+
+---
+// Error: 22-34 expected font family or auto, found boolean
+#set heading(family: lvl => false)
+= Heading
+
+---
+// Error: 22-37 missing argument: b
+#set heading(family: (a, b) => a + b)
+= Heading
+
+---
+// Error: 22-30 unexpected argument
+#set heading(family: () => {})
+= Heading