refactor: add support for custom figure kinds

This also fixes some bugs with referencing figures in elements such as align,
this also generalizes the logic to handle nesting.

#__format-ref show rule now goes through #custom-ref most of the time,
so builtin #ref can be used instead.

Equations are now referenceable and respect numbering per sections.

src/figures.typ

@@ -1,137 +1,302 @@
 // This file includes logic for displaying figures in the document, including numbering, captions, and references. It also includes logic for handling figure numbering per section or globally, based on the configuration set in `config.typ`.
 #import "config.typ": __FIG_PER_SECTION
 
-/// Localizations for various recognized referenceable figure types.
+// ---------------------------------------------------------------------------
-/// This is used to provide the correct supplement text for each type of figure
+// Figure kinds
-/// when generating references, captions and outlines.
+// ---------------------------------------------------------------------------
-#let __figureNames = (
+
-    table: (
+/// Localizations for the recognized referenceable figure kinds, keyed by `repr(func)`.
-        supplement: "lentelė",
+/// - supplement: used in captions (the caption adds the separator, e.g. "1 pav.").
-    ),
+/// - ref-name: used in references, e.g. "(žr. 1 pav.)".
-    image: (
+#let __figure-kinds = (
-        supplement: "pav",
+  table: (supplement: "lentelė", ref-name: "lentelė"),
-    ),
+  image: (supplement: "pav", ref-name: "pav."),
-    raw: (
+  raw: (supplement: "kodo frag", ref-name: "kodo frag."),
-        supplement: "kodo frag",
+  equation: (supplement: "lygtis", ref-name: "formulė"),
-    ),
-    equation: (
-        supplement: "lygtis",
-    ),
 )
 
-/// Helper function to get the supplement text for a given figure type.
+/// Figure kinds registered by the document author via `ktu-paper(figureKinds: ...)`.
-#let __fig-supplement(func) = {
+#let __USER_FIGURE_KINDS = state("ktu-user-figure-kinds", (:))
-    return __figureNames.at(repr(func)).supplement
+
+/// Validates author-supplied kinds and fills in defaults (`ref-name` falls back
+/// to `supplement`).
+#let __normalize-figure-kinds(kinds) = {
+  let normalized = (:)
+  for (key, spec) in kinds {
+    assert(
+      "supplement" in spec,
+      message: "figure kind `" + key + "` needs a `supplement`",
+    )
+    normalized.insert(key, (
+      supplement: spec.supplement,
+      ref-name: spec.at("ref-name", default: spec.supplement),
+    ))
+  }
+  normalized
 }
 
-#let __fig-func-to-name(func) = {
+/// Recognized kinds merged with the author's. Must be called in a context.
-    if func == table {
+#let __all-figure-kinds() = __figure-kinds + __USER_FIGURE_KINDS.get()
-        return "lentelė"
+
-    }
+/// Resets all figure and equation counters (used when numbering per section).
-    if func == image {
+/// Must be called in a context.
-        return "pav."
+#let __reset-figure-counters() = {
-    }
+  let custom = (..__USER_FIGURE_KINDS.get().keys(),)
-    if func == grid {
+  for kind in (image, table, raw, ..custom) {
-        return "pav."
+    counter(figure.where(kind: kind)).update(0)
-    }
+  }
-    return func
+  counter(math.equation).update(0)
 }
 
-/// Helper function to resolve figure numbers (either global "1" or section-based "1.1").
+/// Returns the caption supplement text for a figure kind (e.g. `table`, `math.equation`).
-/// It takes a figure item and referenced location and returns the appropriate number as a string.
+#let __fig-supplement(func) = __figure-kinds.at(repr(func)).supplement
+
+/// Finds the first recognized kind ("table", "image", "raw" or "equation") inside
+/// `node`, descending through wrappers (align, block, grid, stack, ...).
+/// Returns the key into `__figure-kinds`, or `none` if nothing is recognized.
+#let __find-figure-kind(node) = {
+  if type(node) != content { return none }
+
+  let key = repr(node.func())
+  if key in __figure-kinds { return key }
+
+  // Single-child wrappers
+  for field in ("body", "child") {
+    if node.has(field) {
+      let found = __find-figure-kind(node.at(field))
+      if found != none { return found }
+    }
+  }
+
+  // Multi-child containers
+  if node.has("children") {
+    for child in node.children {
+      let found = __find-figure-kind(child)
+      if found != none { return found }
+    }
+  }
+
+  none
+}
+
+/// Resolves the `__figure-kinds` key of a figure element. An explicit custom
+/// (string) kind always wins; otherwise the kind is detected from the body.
+#let __figure-key(fig) = {
+  if type(fig.kind) == str { return fig.kind }
+  let found = __find-figure-kind(fig.body)
+  if found != none { found } else { repr(fig.kind) }
+}
+
+/// Returns the caption supplement of a figure element.
+#let __figure-supplement-of(fig) = {
+  let key = __figure-key(fig)
+  let kinds = __all-figure-kinds()
+  assert(
+    key in kinds,
+    message: "unregistered figure kind `" + key + "`: register it with ktu-paper(figureKinds: ...)",
+  )
+  kinds.at(key).supplement
+}
+
+// ---------------------------------------------------------------------------
+// Numbering
+// ---------------------------------------------------------------------------
+
+/// Resolves a figure's number (global "1" or section-based "1.1") at `loc`,
+/// returned as a string.
 #let __get-figure-number(fig-item, loc) = {
-    if type(fig-item) != content {
+  assert(
-        panic("can only get figure number for content elements")
+    type(fig-item) == content and fig-item.func() == figure,
-    }
+    message: "can only get figure number for figure elements",
+  )
 
-    if fig-item.func() != figure {
+  let fig-counter = counter(figure.where(kind: fig-item.kind))
-        panic("can only get figure number for figure elements")
-    }
 
-    let perSection = __FIG_PER_SECTION.at(loc)
+  if __FIG_PER_SECTION.at(loc) {
-    if perSection {
+    let chapter = counter(heading).at(loc).at(0, default: 0)
-        let chapter = counter(heading).at(loc).at(0, default: 0)
+    let fig-num = fig-counter.at(loc).at(0, default: 1)
-        let figNum = counter(figure.where(kind: fig-item.kind)).at(loc).at(0, default: 1)
+    str(chapter) + "." + str(fig-num)
-        str(chapter) + "." + str(figNum)
+  } else {
-    } else {
+    str(fig-counter.at(loc).at(0))
-        let caption = fig-item.caption
+  }
-        if caption != none and caption.numbering != none {
-            str(caption.counter.at(loc).at(0))
-        } else {
-            str(counter(figure.where(kind: fig-item.kind)).at(loc).at(0))
-        }
-    }
 }
 
-/// Helper function to get the number of a locatable element (figure, table, etc.)
+/// Resolves an equation's number (global "2" or section-based "1.2") at `loc` without parentheses.
-/// based on its location in the document. It returns the number as a string, formatted according to the current numbering scheme.
+/// `loc` is the location of the equation.
+#let __get-equation-number(loc) = {
+  let number = counter(math.equation).at(loc).at(0)
+  if __FIG_PER_SECTION.at(loc) {
+    let chapter = counter(heading).at(loc).at(0, default: 0)
+    str(chapter) + "." + str(number)
+  } else {
+    str(number)
+  }
+}
+
+/// Numbering function for `math.equation`: "(2)" or, when numbering per
+/// section, "(1.2)".
+#let __equation-numbering(..nums) = context {
+  let number = nums.pos().first()
+  if __FIG_PER_SECTION.get() {
+    let chapter = counter(heading).get().at(0, default: 0)
+    // TODO: this is a hack to prevent the numbers from being completely
+    // part of a regex match against math.equation.
+    // It would be better to render math equation numbers as a separate element, so that regexes don't see them at all.
+    [(#box[#chapter].#box[#number])]
+  } else {
+    [(#number)]
+  }
+}
+
+/// Returns the number of a located figure element as a string, formatted
+/// according to the current numbering scheme.
 #let __get-locatable-element-number(element) = context {
-    if type(element) != content {
+  assert(type(element) == content, message: "can only reference elements of type content")
-        panic("can only reference elements of type content")
+  assert(element.func() == figure, message: "referenced element does not have a figure function")
-    }
+  __get-figure-number(element, element.location())
-    let loc = element.location()
-    let elem = query(loc)
-    let fig = elem.at(0)
-    if fig.func() != figure {
-        panic("referenced element does not have a figure function")
-    }
-    __get-figure-number(fig, loc)
 }
 
-/// Formats a reference to a figure, table, or other numbered element.
+/// Appendix headings carry an "appendix-heading" metadata marker as their
-#let __format-ref(it) = {
+/// first child and are numbered differently from normal headings.
-    let el = it.element
+#let __is-appendix-heading(item) = {
-    if el == none {
+  let body = item.body
-        return it
+  body.has("children") and body.children.at(0, default: none) == metadata("appendix-heading")
-    }
+}
 
-    // Override equation references
+/// Parses a dotted number such as "12" or "1.12.5" into an array of ints.
-    if el.func() == math.equation {
+/// Returns `none` if any part is not purely numeric.
-        link(
+#let __parse-number(value) = {
-            el.location(),
+  let parts = value.split(".")
-            numbering(el.numbering, ..counter(math.equation).at(el.location()))
+  if parts.all(p => p.match(regex("^\\d+$")) != none) {
-        )
+    parts.map(int)
-    // Override references to figures
+  } else {
-    } else if el.func() == figure {
+    none
-        let fig = el.body
+  }
-        let figType = fig.func()
+}
-        let numStr = __get-figure-number(el, el.location())
+
-        link(el.location(), [
+/// Sorts `(location, number)` entries. If every number is numeric (optionally
-            (žr. #numStr #__fig-func-to-name(figType))
+/// dotted, e.g. "1.12.5"), sorts component-wise numerically; otherwise falls
-        ])
+/// back to plain string ordering.
+#let __sort-entries(entries) = {
+  let keys = entries.map(e => __parse-number(e.at(1)))
+  if keys.all(k => k != none) {
+    entries.zip(keys).sorted(key: pair => pair.at(1)).map(pair => pair.at(0))
+  } else {
+    entries.sorted(key: e => e.at(1))
+  }
+}
+
+// ---------------------------------------------------------------------------
+// References
+// ---------------------------------------------------------------------------
+
+/// Returns a number for the given reference target.
+/// - target (label): Label to get the reference number for.
+#let ref-no(target) = context {
+  if type(target) != label {
+    panic("ref-no expects a label as input")
+  }
+
+  let item = query(target).first()
+  let loc = item.location()
+  let kind = item.func()
+
+  let number = if kind == figure {
+    __get-figure-number(item, loc)
+  } else if kind == math.equation {
+    "(" + __get-equation-number(loc) + ")"
+  } else if kind == heading {
+    let numbers = counter(heading).at(loc)
+    if __is-appendix-heading(item) {
+      str(numbers.at(1))
     } else {
-        it
+      numbers.map(str).join(".")
     }
+  } else {
+    panic("referencing unknown type: " + repr(kind))
+  }
+
+  link(loc, number)
 }
 
+/// Create a custom reference for multiple targets of the same type.
+/// - targets (label...): Labels to create a custom reference for.
+#let custom-ref(..targets) = context {
+  let group-name = none
+  let entries = () // (location, number) pairs
+
+  for target in targets.pos() {
+    let item = query(target).first()
+    let loc = item.location()
+
+    let (name, number) = if item.func() == math.equation {
+      (__figure-kinds.equation.ref-name, __get-equation-number(loc))
+    } else if item.func() == figure {
+      let kind = __figure-key(item)
+      let kinds = __all-figure-kinds()
+      if kind not in kinds {
+        panic("unregistered figure kind `" + kind + "`: register it with ktu-paper(figureKinds: ...)")
+      }
+      (kinds.at(kind).ref-name, __get-figure-number(item, loc))
+    } else {
+      panic("referencing unknown type: " + repr(item.func()))
+    }
+
+    if group-name != none and group-name != name {
+      panic("custom-ref does not support grouping different elements together")
+    }
+    group-name = name
+    entries.push((loc, number))
+  }
+
+  if entries.len() > 0 {
+    let links = __sort-entries(entries).map(e => link(e.at(0), e.at(1)))
+    [(žr. #links.join(", ") #group-name)]
+  }
+}
+
+/// Show rule for `@label` references. Equations render as a bare "(1.2)";
+/// figures delegate to `custom-ref`, so "(žr. 1 pav.)" is built in one place.
+/// Anything else (headings, citations, unregistered figure kinds) keeps
+/// Typst's default rendering.
+#let __format-ref(it) = {
+  let el = it.element
+  if el == none { return it }
+
+  if el.func() == math.equation {
+    let loc = el.location()
+    link(loc, "(" + __get-equation-number(loc) + ")")
+  } else if el.func() == figure and __figure-key(el) in __all-figure-kinds() {
+    custom-ref(it.target)
+  } else {
+    it
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Figure layout
+// ---------------------------------------------------------------------------
+
 /// Formats a figure with a caption and number.
 #let __format-figure(it) = {
-    let separator = [.]
+  let caption = it.caption
-    let caption = it.caption
 
-    // If caption is not provided, just show the figure body
+  // If caption is not provided, just show the figure body
-    if caption == none {
+  if caption == none { return it.body }
-        return it.body
-    }
 
-    let counter = caption.counter  
+  set text(size: 11pt)
-    let supplement = caption.supplement
 
-    if it.body.func() == math.equation {
+  let supplement = __figure-supplement-of(it)
-        supplement = __fig-supplement("equation")
+  let separator = [.]
-    } else {
+  let fig-num = context __get-figure-number(it, here())
-        supplement = __fig-supplement(it.kind)
+  let caption-line = [*#fig-num #supplement#separator* #caption.body]
-    }
-    set text(size: 11pt)
 
-    let fig-num = context __get-figure-number(it, here())
+  // Tables have captions at the top
-
+  if it.kind == table {
-    // Tables have captions at the top
+    align(left, caption-line)
-    if it.kind == table {
+    it.body
-        align(left)[*#fig-num #supplement#separator* #caption.body]
+  } else {
-        it.body
+    it.body
-    } else {
+    align(center, caption-line)
-        it.body
+  }
-        align(center)[*#fig-num #supplement#separator* #caption.body]
-    }
 }

src/lib.typ

@@ -10,10 +10,15 @@
     ktu-academic-honestly-declaration-page
 )
 #import "figures.typ": (
-    __fig-supplement,
+    ref-no, custom-ref,
 
-    __get-figure-number,
+    __fig-supplement,
+    __figure-supplement-of,
     __get-locatable-element-number,
+    __reset-figure-counters,
+    __equation-numbering,
+    __normalize-figure-kinds,
+    __USER_FIGURE_KINDS,
 
     __format-ref,
     __format-figure,
@@ -85,109 +90,6 @@
   integer-part + padded-decimal
 }
 
-/// Returns a number for the given reference target.
-/// - target (label): Label to get the reference number for.
-#let ref-no(target) = context {
-    if type(target) != label {
-        panic("ref-no expects a label as input")
-    }
-
-    let loc = locate(target)
-    let elem = query(target)
-    let item = elem.at(0)
-    
-    // If it's a figure, return the figure number
-    if item.func() == figure {
-        link(loc, [#__get-figure-number(item, loc)])
-    } else if item.func() == heading {
-        let fullValue = counter(heading).at(loc)
-
-        // If the heading has the "appendix-heading" metadata attached,
-        // it means that it's an appendix heading and should be numbered differently
-        if item.body.has("children") and item.body.children.at(0) == metadata("appendix-heading") {
-            link(loc, [#fullValue.at(1)])
-        // Otherwise, show the fully qualified heading numbering
-        } else {
-            link(loc, [#fullValue.map((x) => str(x)).join(".")])
-        }
-    } else {
-        panic("referencing unknown type: ", item.func())
-    }
-}
-
-// Constructs a reference from multiple targets
-#let custom-ref(..targets) = context {
-  let collected = (:)
-
-  for target in targets.pos() {
-    let loc = locate(target)
-    let elem = query(target)
-
-    let fig = elem.at(0)
-    if fig.func() != figure {
-      panic("referencing unknown type")
-    }
-
-    // determine what we...
-    let element = fig.body;
-    let elementName = "unknown";
-    let found = false;
-    while not found {
-
-      // If it's a grid
-      if (element.func() == grid) {
-        for cell in element.children {
-          if cell.body.func() == image {
-            elementName = "pav."
-            found = true
-            break
-          }
-        }
-        break
-      }
-
-      if (element.func() == table) {
-        elementName = "lentelė"
-        found = true
-        break
-      }
-
-    if (element.func() == raw) {
-        elementName = "kodo frag."
-        found = true
-        break
-      }
-
-      if (element.func() == image) {
-        elementName = "pav."
-        found = true
-        break
-      }
-
-      panic("unknown element", element.func())
-    }
-
-    if elementName not in collected {
-      collected.insert(elementName, ())
-    }
-
-    collected.at(elementName).push(
-      (loc, __get-figure-number(fig, loc))
-    )
-  }
-
-  if collected.len() > 1 {
-    panic("multiple types detected")
-  }
-
-  for type in collected {
-    let name = type.at(0)
-    let arr = type.at(1).sorted(key: k => k.at(1))
-
-    [(žr. #arr.map(x =>link(x.at(0), [#x.at(1)])).join(", ") #name)]
-  }
-}
-
 
 /// Returns a KTU template-conformant outline.
 /// - depth (int | none): Same as outline depth.
@@ -297,11 +199,12 @@
                 #(o.target == selector(heading))
             ]*/
 
+            // Any other figure kind (generally, user defined kinds)
             show outline.entry: it => link(
                 it.element.location(),
                 it.indented(
-                    [*#__get-locatable-element-number(it.element) unknown.*],
+                    [*#__get-locatable-element-number(it.element) #__figure-supplement-of(it.element).*],
-                    it.inner(), gap: 0pt
+                    [ #it.inner()], gap: 0pt
                 ),
             )
             o
@@ -342,6 +245,10 @@
     ktu-outline(title: heading("Paveikslų sąrašas", outlined: true), target: figure.where(kind: image))
 }
 
+#let ktu-code-list() = {
+  ktu-outline(title: heading("Kodo fragmentų sąrašas", outlined: true), target: figure.where(kind: raw))
+}
+
 /// Define a custom outline for a specific figure type.
 /// This function creates an outline entry with a specified title and figure type.
 /// The outline entry will link to the figure of the specified type.
@@ -389,7 +296,7 @@
     set heading(numbering: "1.")
 
     // Number equations
-    set math.equation(numbering: "(1)")
+    set math.equation(numbering: __equation-numbering)
 
     // Set text font
     set text(
@@ -480,9 +387,7 @@
 
             // Reset figure counters if numbering per section is enabled
             if __FIG_PER_SECTION.get() {
-                counter(figure.where(kind: image)).update(0)
+                __reset-figure-counters()
-                counter(figure.where(kind: table)).update(0)
-                counter(figure.where(kind: raw)).update(0)
             }
 
             // po antraštės - 10 pt
@@ -536,9 +441,13 @@
 /// - figureNumberingPerSection (bool): Whether to number figures per section rather than globally.
 /// - uppercaseSectionTitles (bool): Whether to uppercase section titles.
 /// - unnumberedHeadingAlignment (alignment): Alignment for unnumbered headings.
+/// - figureKinds (dictionary): Extra figure kinds, e.g.
+///   `(command: (supplement: "komanda", ref-name: "komanda"))`. `ref-name`
+///   defaults to `supplement`. Use them with `figure(kind: "command", ...)`.
 #let ktu-paper(
     font: "Times New Roman",
     figureNumberingPerSection: false,
+    figureKinds: (:),
     // Not sure about formal conformity of deviating
     // from the defaults of the next options,
     // but some modules may require them
@@ -552,6 +461,7 @@
         unnumberedHeadingAlignment: unnumberedHeadingAlignment,
     )
     __FIG_PER_SECTION.update(figureNumberingPerSection)
+    __USER_FIGURE_KINDS.update(__normalize-figure-kinds(figureKinds))
     __UPPERCASE_SECTION_TITLES.update(uppercaseSectionTitles)
     body
 }