use std::fmt::Write;
use std::num::NonZeroUsize;
use std::str::FromStr;
use std::sync::Arc;
use comemo::{Track, Tracked, TrackedMut};
use ecow::{EcoString, EcoVec, eco_format, eco_vec};
use smallvec::{SmallVec, smallvec};
use typst_syntax::Span;
use typst_utils::{LazyHash, NonZeroExt, Protected};
use crate::diag::{At, HintedStrResult, SourceDiagnostic, SourceResult, bail, warning};
use crate::engine::{Engine, Route, Sink, Traced};
use crate::foundations::{
Args, Array, Construct, Content, Context, Element, Func, IntoValue, Label,
LocatableSelector, NativeElement, Packed, Repr, Selector, ShowFn, Smart, Str,
StyleChain, Value, cast, elem, func, scope, select_where, ty,
};
use crate::introspection::{
History, Introspect, Introspector, Locatable, Location, QueryFirstIntrospection, Tag,
Unqueriable,
};
use crate::layout::{Frame, FrameItem, PageElem};
use crate::math::EquationElem;
use crate::model::{FigureElem, FootnoteElem, HeadingElem, Numbering, NumberingPattern};
use crate::{Library, World};
#[ty(scope)]
#[derive(Debug, Clone, PartialEq, Hash)]
pub struct Counter(CounterKey);
impl Counter {
pub fn new(key: CounterKey) -> Counter {
Self(key)
}
pub fn of(func: Element) -> Self {
Self::new(CounterKey::Selector(Selector::Elem(func, None)))
}
pub fn select_any() -> Selector {
CounterUpdateElem::ELEM.select()
}
pub fn select(
&self,
introspector: Tracked<dyn Introspector + '_>,
loc: Location,
) -> Selector {
let mut selector = select_where!(CounterUpdateElem, key => self.0.clone());
match &self.0 {
CounterKey::Page => {
if let Some(doc_location) = introspector.document(loc) {
selector = Selector::Within {
selector: Arc::new(selector),
ancestor: Arc::new(doc_location.into()),
};
}
}
CounterKey::Selector(key) => {
selector = Selector::Or(eco_vec![selector, key.clone()]);
}
CounterKey::Str(_) => {}
}
selector
}
pub fn is_page(&self) -> bool {
self.0 == CounterKey::Page
}
pub fn display_at(
&self,
engine: &mut Engine,
loc: Location,
styles: StyleChain,
numbering: &Numbering,
span: Span,
) -> SourceResult<Content> {
let context = Context::new(Some(loc), Some(styles));
Ok(engine
.introspect(CounterAtIntrospection(self.clone(), loc, span))?
.display(engine, context.track(), span, numbering)?
.display())
}
fn matching_numbering(
&self,
engine: &mut Engine,
styles: StyleChain,
loc: Location,
span: Span,
) -> Option<Numbering> {
match self.0 {
CounterKey::Page => loc.page_numbering(engine, span),
CounterKey::Selector(Selector::Elem(func, _)) => engine
.introspect(QueryFirstIntrospection(Selector::Location(loc), span))
.and_then(|content| {
if func == HeadingElem::ELEM {
content
.to_packed::<HeadingElem>()
.and_then(|elem| elem.numbering.as_option().clone())
.flatten()
} else if func == FigureElem::ELEM {
content
.to_packed::<FigureElem>()
.and_then(|elem| elem.numbering.as_option().clone())
.flatten()
} else if func == EquationElem::ELEM {
content
.to_packed::<EquationElem>()
.and_then(|elem| elem.numbering.as_option().clone())
.flatten()
} else if func == FootnoteElem::ELEM {
content
.to_packed::<FootnoteElem>()
.and_then(|elem| elem.numbering.as_option().clone())
} else {
None
}
})
.or_else(|| {
if func == HeadingElem::ELEM {
styles.get_cloned(HeadingElem::numbering)
} else if func == FigureElem::ELEM {
styles.get_cloned(FigureElem::numbering)
} else if func == EquationElem::ELEM {
styles.get_cloned(EquationElem::numbering)
} else if func == FootnoteElem::ELEM {
Some(styles.get_cloned(FootnoteElem::numbering))
} else {
None
}
}),
_ => None,
}
}
}
#[scope]
impl Counter {
#[func(constructor)]
pub fn construct(
key: CounterKey,
) -> Counter {
Self::new(key)
}
#[func(contextual)]
pub fn get(
&self,
engine: &mut Engine,
context: Tracked<Context>,
span: Span,
) -> SourceResult<CounterState> {
let loc = context.location().at(span)?;
engine.introspect(CounterAtIntrospection(self.clone(), loc, span))
}
#[func(contextual)]
pub fn display(
self,
engine: &mut Engine,
context: Tracked<Context>,
span: Span,
#[default]
numbering: Smart<Numbering>,
#[named]
#[default]
at: Smart<LocatableSelector>,
#[named]
#[default(false)]
both: bool,
) -> SourceResult<Value> {
let location = match at {
Smart::Auto => context.location().at(span)?,
Smart::Custom(ref selector) => {
selector.resolve_unique(engine, context, span)?
}
};
let state = if both {
engine.introspect(CounterBothIntrospection(self.clone(), location, span))?
} else {
engine.introspect(CounterAtIntrospection(self.clone(), location, span))?
};
let numbering = numbering
.custom()
.or_else(|| {
self.matching_numbering(engine, context.styles().ok()?, location, span)
})
.unwrap_or_else(|| NumberingPattern::from_str("1.1").unwrap().into());
if at.is_custom() {
let context = Context::new(Some(location), context.styles().ok());
state.display(engine, context.track(), span, &numbering)
} else {
state.display(engine, context, span, &numbering)
}
}
#[func(contextual)]
pub fn at(
&self,
engine: &mut Engine,
context: Tracked<Context>,
span: Span,
selector: LocatableSelector,
) -> SourceResult<CounterState> {
let loc = selector.resolve_unique(engine, context, span)?;
engine.introspect(CounterAtIntrospection(self.clone(), loc, span))
}
#[func(contextual)]
pub fn final_(
&self,
engine: &mut Engine,
context: Tracked<Context>,
span: Span,
) -> SourceResult<CounterState> {
let loc = context.location().at(span)?;
engine.introspect(CounterFinalIntrospection(self.clone(), loc, span))
}
#[func]
pub fn step(
self,
span: Span,
#[named]
#[default(NonZeroUsize::ONE)]
level: NonZeroUsize,
) -> Content {
self.update(span, CounterUpdate::Step(level))
}
#[func]
pub fn update(
self,
span: Span,
update: CounterUpdate,
) -> Content {
CounterUpdateElem::new(self.0, update).pack().spanned(span)
}
}
impl Repr for Counter {
fn repr(&self) -> EcoString {
eco_format!("counter({})", self.0.repr())
}
}
#[derive(Debug, Clone, PartialEq, Hash)]
pub enum CounterKey {
Page,
Selector(Selector),
Str(Str),
}
cast! {
CounterKey,
self => match self {
Self::Page => PageElem::ELEM.into_value(),
Self::Selector(v) => v.into_value(),
Self::Str(v) => v.into_value(),
},
v: Str => Self::Str(v),
v: Label => Self::Selector(Selector::Label(v)),
v: Element => {
if v == PageElem::ELEM {
Self::Page
} else {
Self::Selector(LocatableSelector::from_value(v.into_value())?.0)
}
},
v: LocatableSelector => Self::Selector(v.0),
}
impl Repr for CounterKey {
fn repr(&self) -> EcoString {
match self {
Self::Page => "page".into(),
Self::Selector(selector) => selector.repr(),
Self::Str(str) => str.repr(),
}
}
}
#[derive(Debug, Clone, PartialEq, Hash)]
pub enum CounterUpdate {
Set(CounterState),
Step(NonZeroUsize),
Func(Func),
}
cast! {
CounterUpdate,
v: CounterState => Self::Set(v),
v: Func => Self::Func(v),
}
pub trait Count {
fn update(&self) -> Option<CounterUpdate>;
}
#[derive(Debug, Clone, PartialEq, Hash)]
pub struct CounterState(pub SmallVec<[u64; 3]>);
impl CounterState {
pub fn init(page: bool) -> Self {
Self(smallvec![u64::from(page)])
}
pub fn update(
&mut self,
engine: &mut Engine,
update: CounterUpdate,
) -> SourceResult<()> {
match update {
CounterUpdate::Set(state) => *self = state,
CounterUpdate::Step(level) => self.step(level, 1),
CounterUpdate::Func(func) => {
*self = func
.call(engine, Context::none().track(), self.0.iter().copied())?
.cast()
.at(func.span())?
}
}
Ok(())
}
pub fn step(&mut self, level: NonZeroUsize, by: u64) {
let level = level.get();
while self.0.len() < level {
self.0.push(0);
}
self.0[level - 1] = self.0[level - 1].saturating_add(by);
self.0.truncate(level);
}
pub fn first(&self) -> u64 {
self.0.first().copied().unwrap_or(1)
}
pub fn display(
&self,
engine: &mut Engine,
context: Tracked<Context>,
span: Span,
numbering: &Numbering,
) -> SourceResult<Value> {
numbering.apply(engine, context, span, &self.0)
}
}
cast! {
CounterState,
self => Value::Array(self.0.into_iter().map(IntoValue::into_value).collect()),
num: u64 => Self(smallvec![num]),
array: Array => Self(array
.into_iter()
.map(Value::cast)
.collect::<HintedStrResult<_>>()?),
}
#[elem(Construct, Locatable, Count)]
pub struct CounterUpdateElem {
#[required]
key: CounterKey,
#[required]
#[internal]
update: CounterUpdate,
}
impl Construct for CounterUpdateElem {
fn construct(_: &mut Engine, args: &mut Args) -> SourceResult<Content> {
bail!(args.span, "cannot be constructed manually");
}
}
impl Count for Packed<CounterUpdateElem> {
fn update(&self) -> Option<CounterUpdate> {
Some(self.update.clone())
}
}
#[elem(Construct, Unqueriable, Locatable)]
pub struct CounterDisplayElem {
#[required]
#[internal]
counter: Counter,
#[required]
#[internal]
numbering: Smart<Numbering>,
#[required]
#[internal]
both: bool,
}
impl Construct for CounterDisplayElem {
fn construct(_: &mut Engine, args: &mut Args) -> SourceResult<Content> {
bail!(args.span, "cannot be constructed manually");
}
}
pub const COUNTER_DISPLAY_RULE: ShowFn<CounterDisplayElem> = |elem, engine, styles| {
Ok(elem
.counter
.clone()
.display(
engine,
Context::new(elem.location(), Some(styles)).track(),
elem.span(),
elem.numbering.clone(),
Smart::Auto,
elem.both,
)?
.display())
};
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub struct ManualPageCounter {
physical: NonZeroUsize,
logical: u64,
}
impl ManualPageCounter {
pub fn new() -> Self {
Self { physical: NonZeroUsize::ONE, logical: 1 }
}
pub fn physical(&self) -> NonZeroUsize {
self.physical
}
pub fn logical(&self) -> u64 {
self.logical
}
pub fn visit(&mut self, engine: &mut Engine, page: &Frame) -> SourceResult<()> {
for (_, item) in page.items() {
match item {
FrameItem::Group(group) => self.visit(engine, &group.frame)?,
FrameItem::Tag(Tag::Start(elem, _)) => {
let Some(elem) = elem.to_packed::<CounterUpdateElem>() else {
continue;
};
if elem.key == CounterKey::Page {
let mut state = CounterState(smallvec![self.logical]);
state.update(engine, elem.update.clone())?;
self.logical = state.first();
}
}
_ => {}
}
}
Ok(())
}
pub fn step(&mut self) {
self.physical = self.physical.saturating_add(1);
self.logical += 1;
}
}
impl Default for ManualPageCounter {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, PartialEq, Hash)]
struct CounterAtIntrospection(Counter, Location, Span);
impl Introspect for CounterAtIntrospection {
type Output = SourceResult<CounterState>;
fn introspect(
&self,
engine: &mut Engine,
introspector: Tracked<dyn Introspector + '_>,
) -> Self::Output {
let Self(counter, loc, _) = self;
let selector = counter.select(introspector, *loc);
let sequence = sequence(counter, &selector, engine, introspector)?;
let offset = introspector.query_count_before(&selector, *loc);
let (mut state, page) = sequence[offset].clone();
if counter.is_page() {
let delta = introspector
.page(*loc)
.unwrap_or(NonZeroUsize::ONE)
.get()
.saturating_sub(page.get());
state.step(NonZeroUsize::ONE, delta as u64);
}
Ok(state)
}
fn diagnose(&self, history: &History<Self::Output>) -> SourceDiagnostic {
format_convergence_warning(&self.0, self.2, history)
}
}
#[derive(Debug, Clone, PartialEq, Hash)]
struct CounterBothIntrospection(Counter, Location, Span);
impl Introspect for CounterBothIntrospection {
type Output = SourceResult<CounterState>;
fn introspect(
&self,
engine: &mut Engine,
introspector: Tracked<dyn Introspector + '_>,
) -> Self::Output {
let Self(counter, loc, _) = self;
let selector = counter.select(introspector, *loc);
let sequence = sequence(counter, &selector, engine, introspector)?;
let offset = introspector.query_count_before(&selector, *loc);
let (mut at_state, at_page) = sequence[offset].clone();
let (mut final_state, final_page) = sequence.last().unwrap().clone();
if counter.is_page() {
let at_delta = introspector
.page(*loc)
.unwrap_or(NonZeroUsize::ONE)
.get()
.saturating_sub(at_page.get());
at_state.step(NonZeroUsize::ONE, at_delta as u64);
let final_delta = introspector
.pages(*loc)
.unwrap_or(NonZeroUsize::ONE)
.get()
.saturating_sub(final_page.get());
final_state.step(NonZeroUsize::ONE, final_delta as u64);
}
Ok(CounterState(smallvec![at_state.first(), final_state.first()]))
}
fn diagnose(&self, history: &History<Self::Output>) -> SourceDiagnostic {
format_convergence_warning(&self.0, self.2, history)
}
}
#[derive(Debug, Clone, PartialEq, Hash)]
struct CounterFinalIntrospection(Counter, Location, Span);
impl Introspect for CounterFinalIntrospection {
type Output = SourceResult<CounterState>;
fn introspect(
&self,
engine: &mut Engine,
introspector: Tracked<dyn Introspector + '_>,
) -> Self::Output {
let Self(counter, loc, _) = self;
let selector = counter.select(introspector, *loc);
let sequence = sequence(counter, &selector, engine, introspector)?;
let (mut state, page) = sequence.last().unwrap().clone();
if counter.is_page() {
let delta = introspector
.pages(*loc)
.unwrap_or(NonZeroUsize::ONE)
.get()
.saturating_sub(page.get());
state.step(NonZeroUsize::ONE, delta as u64);
}
Ok(state)
}
fn diagnose(&self, history: &History<Self::Output>) -> SourceDiagnostic {
format_convergence_warning(&self.0, self.2, history)
}
}
fn sequence(
counter: &Counter,
selector: &Selector,
engine: &mut Engine,
introspector: Tracked<dyn Introspector + '_>,
) -> SourceResult<EcoVec<(CounterState, NonZeroUsize)>> {
sequence_impl(
counter,
selector,
engine.world,
engine.library,
introspector,
engine.traced,
TrackedMut::reborrow_mut(&mut engine.sink),
engine.route.track(),
)
}
#[comemo::memoize]
#[allow(clippy::too_many_arguments)]
fn sequence_impl(
counter: &Counter,
selector: &Selector,
world: Tracked<dyn World + '_>,
library: &LazyHash<Library>,
introspector: Tracked<dyn Introspector + '_>,
traced: Tracked<Traced>,
sink: TrackedMut<Sink>,
route: Tracked<Route>,
) -> SourceResult<EcoVec<(CounterState, NonZeroUsize)>> {
let mut engine = Engine {
library,
world,
introspector: Protected::from_raw(introspector),
traced,
sink,
route: Route::extend(route).unnested(),
};
let mut current = CounterState::init(matches!(counter.0, CounterKey::Page));
let mut page = NonZeroUsize::ONE;
let mut stops = eco_vec![(current.clone(), page)];
for elem in introspector.query(selector) {
if counter.is_page() {
let prev = page;
page = introspector
.page(elem.location().unwrap())
.unwrap_or(NonZeroUsize::ONE);
let delta = page.get() - prev.get();
if delta > 0 {
current.step(NonZeroUsize::ONE, delta as u64);
}
}
if let Some(update) = match elem.with::<dyn Count>() {
Some(countable) => countable.update(),
None => Some(CounterUpdate::Step(NonZeroUsize::ONE)),
} {
current.update(&mut engine, update)?;
}
stops.push((current.clone(), page));
}
Ok(stops)
}
fn format_convergence_warning(
counter: &Counter,
span: Span,
history: &History<SourceResult<CounterState>>,
) -> SourceDiagnostic {
warning!(span, "value of {} did not converge", format_key(counter)).with_hint(
history.hint("values", |ret| match ret {
Ok(v) => format_counter_state(v),
Err(_) => "(errored)".into(),
}),
)
}
fn format_key(counter: &Counter) -> EcoString {
match counter.0 {
CounterKey::Page => "the page counter".into(),
_ => eco_format!("`{}`", counter.repr()),
}
}
fn format_counter_state(state: &CounterState) -> EcoString {
let mut output = EcoString::new();
let mut sep = "";
for value in state.0.as_slice() {
write!(output, "{sep}{value}").unwrap();
sep = ", ";
}
output
}