-- lualineno version = 0.1, 2025-11-28

-- \seccc Initialization^^M
-- Some declarations of local functions/constants of global ones
-- to avoid table lookups.

local runtoks = tex.runtoks
local put_next = token.unchecked_put_next
local create = token.create
local new_tok = token.new
local lbrace = new_tok(string.byte('{'), token.command_id'left_brace')
local rbrace = new_tok(string.byte('}'), token.command_id'right_brace')
local b = new_tok(string.byte('b'), token.command_id'other_char')
local d = new_tok(string.byte('d'), token.command_id'other_char')
local i = new_tok(string.byte('i'), token.command_id'other_char')
local r = new_tok(string.byte('r'), token.command_id'other_char')
local zero = new_tok(string.byte('0'), token.command_id'other_char')
local get_next = token.get_next
local scan_toks = token.scan_toks
local scan_string = token.scan_string
local scan_list = token.scan_list
local scan_int = token.scan_int

-- Sadly there isn't a nice way in \LuaTeX/ to get 
-- a primitive token without using a csname.
-- To be sure `\hbox` has the correct meaning
-- we can use `tex.enableprimitives` to create
-- a new csname with the meaning of the primitive,
-- then create a token with the same `.mode` and `.command`
-- fields so we won't need the csname anymore.
-- All of this is to avoid to use some implementation details
-- (`local hbox = new_tok(141, 21)`)

local hbox do
  local prefix = '@lua^line&no_'
  while token.is_defined(prefix .. 'hbox') do
    prefix = prefix .. '@lua^line&no_'
  end
  tex.enableprimitives(prefix,{'hbox'})
  local tok = create(prefix .. 'hbox')
  hbox  = new_tok(tok.mode, tok.command)
end

local hlist_id = node.id('hlist')
local vlist_id = node.id('vlist')
local glyph_id = node.id('glyph')
local whatsit_id = node.id('whatsit')
local node_write = node.write
local tail = node.tail
local get_attribute = node.get_attribute
local set_attribute = node.set_attribute
local node_flush = node.flush_node
local insert_before = node.insert_before
local insert_after = node.insert_after
local traverse = node.traverse
local rangedimensions = node.rangedimensions
local node_copy = node.copy
local base_kern = node.new('kern', 'user')
local line_sub, eq_sub, align_sub, box_sub
local ignored_subtypes = {}
for k,v in pairs(node.subtypes("hlist")) do
    if v == "line" then line_sub = k end
    if v == "alignment" then align_sub = k end
    if v == "box" then box_sub = k end
    if v == "equation" then eq_sub = k end
    if v == "equationnumber" then ignored_subtypes[k] = true end
    if v == "mathchar" then ignored_subtypes[k] = true end
end
local displayalign_sub = #node.subtypes("hlist") + 1
for k,v in pairs(node.subtypes("vlist")) do
    if v == "vextensible" then ignored_subtypes[k] = true end
    if v == "vdelimiter" then ignored_subtypes[k] = true end
end

local setattribute = tex.setattribute
local texerror = tex.error
local texnest = tex.nest
local format = tex.formatname

-- The module currently works with
-- \OpTeX, \LaTeX/ or Plain.

local optex, latex, plain
if format:find("optex") then
    optex = true
elseif format:find("latex") then
    latex = true
elseif format == "luatex" or
       format == "luahbtex" or
       format:find("plain")
then
    plain = true
end
if not (optex or latex or plain) then
    error("lualineno: The format " .. format .. " is not supported\n\n" ..
             "Use OpTeX, LuaLaTeX or Plain.")
end

local lineno_types = { }
local lineno_attrs = { }
local LINENO_NUMBER = 0x1
local LINENO_RECURSE = 0x2
local lineno_marks = {
    anchor = 1,
    processed = 2,
    display = 3,
}
local type_attr = luatexbase and luatexbase.new_attribute('lualineno_type') or 0
local mark_attr = luatexbase and luatexbase.new_attribute('lualineno_mark') or 1
-- Labels are carried by a user defined whatsit (see the labels section).
-- A real allocator is used when available, otherwise the id defaults to 0;
-- it can be reassigned with the `whatsitid` key, like the attributes.
local label_whatsit = luatexbase and luatexbase.new_whatsit
    and luatexbase.new_whatsit('lualineno') or 0
local user_defined_sub = node.subtype("user_defined")
local base_whatsit = node.new('whatsit', user_defined_sub)
base_whatsit.user_id = label_whatsit
base_whatsit.type = string.byte('l')

local unset_attr = -0x7FFFFFFF

-- We use the `luakeyval` module for the user interface

local keyval = require('luakeyval')
local scan_choice = keyval.choices
local scan_bool = keyval.bool
local process_keys = keyval.process
local messages = {
    error1 = "lualineno: Wrong syntax in \\lualineno",
    value_forbidden = 'lualineno: The key "%s" does not accept a value',
    value_required = 'lualineno: The key "%s" requires a value',
}

-- \seccc Numbering Lines^^M
-- In here we define the main functions of the module,
-- the functions that find and number the lines in a page.
--
-- The following function is used to number a line that is considered
-- \"real" (i.e. has some glyphs in it that are not equation number or 
-- a big math delimiter). This function is used in the `lualineno.numbering`
-- callback, so it can be replaced if desired.
--
-- `line` is the hlist node representing the line, `line_type` is a lua table
-- with the parameters defined in the define key according to the attribute
-- and the column, `offset` is the total shift calculated from the start of
-- the line and `width` is the width of the column containing the lines.

local function number_line(line, line_type, offset, width, dir)

   local head = line.head
   local is_offset = line_type['offset']
   
-- In case \LaTeX/ is used without the luacolor package,
-- we add an additional group to make the boxes color safe.
-- Since `token.scan_list` is ran in horizontal mode, the 
-- default box direction is `\textdirection`, and this can be
-- pretty random at shipout time, so an explicit direction is specified.
-- Currently LTR is used, if someone ask an option can be added, but
-- you can always use `\hbox bidr 1 {}` inside the box.

    put_next({rbrace, rbrace})
    put_next(line_type.left)
    put_next({hbox,b,d,i,r,zero,lbrace,lbrace})
    put_next({rbrace, rbrace})
    put_next(line_type.right)
    put_next({hbox,b,d,i,r,zero,lbrace,lbrace})
    
-- To make sure \"right" always means right, we check the line direction.

    local end_box, start_box  
    if line.dir == "TLT" then
        end_box = scan_list()
        start_box = scan_list()
    else
        start_box = scan_list()
        end_box = scan_list()
    end
    
-- If the vertical list containing the line and the line has different directions
-- we need to mirror the kerns as the kern means opposite direction then the shift,
-- or alignment (similar to how `\shapemode` is working). 
    
    local start_kern_width = 0
    local end_kern_width = 0
    if is_offset then 
        if line.dir == dir then
            start_kern_width = offset
            end_kern_width = width - line.width - offset
        else
            start_kern_width = width - line.width - offset
            end_kern_width = offset
        end
    end
    
    if start_box.head then
        if start_kern_width ~= 0 then
            local start_kern = node_copy(base_kern)
            start_kern.kern = start_kern_width
            head = insert_before(head,head,start_kern)
        end
        head = insert_before(head,head,start_box)
        start_kern_width = -start_box.width - start_kern_width
        if start_kern_width ~= 0 then
            local start_kern = node_copy(base_kern)
            start_kern.kern = start_kern_width
            head = insert_before(head,head,start_kern)
        end
    else
        node_flush(start_box)
    end
    if end_box.head then
        if end_kern_width ~= 0 then
            local end_kern = node_copy(base_kern)
            end_kern.kern = end_kern_width
            head = insert_after(head,tail(head),end_kern)
        end
        head = insert_after(head,tail(head),end_box)
    else
        node_flush(end_box)
    end
    line.head = head
end

local lineno_callbacks
if luatexbase then
    luatexbase.create_callback('lualineno.pre_numbering', 'simple', false)
    luatexbase.create_callback('lualineno.numbering', 'exclusive', number_line)
    luatexbase.create_callback('lualineno.post_numbering', 'simple', false)
    luatexbase.add_to_callback('lualineno.pre_numbering', function(_, line_type)
        runtoks(function() put_next(line_type['toks']) end) 
    end, 'lualineno.runtoks')
    local call_callback = luatexbase.call_callback
    lineno_callbacks = function(line, line_type, offset, width, dir)
        call_callback('lualineno.pre_numbering', line, line_type, offset, width, dir)
        call_callback('lualineno.numbering', line, line_type, offset, width, dir)
        call_callback('lualineno.post_numbering', line, line_type, offset, width, dir)
    end
else
    lineno_callbacks = function(line, line_type, offset, width, dir)
        runtoks(function() put_next(line_type['toks']) end)
        number_line(line, line_type, offset, width, dir)
    end
end

-- Not every object that would be considered a line from \LuaTeX's point of view
-- would be considered a line from a human perspective. For example, a line
-- containing only an indent box, or an alignment containing only rules,
-- so we use the following two functions to search for a glyph node recursively,
-- while ignoring boxes
-- for equation number, for big delimiters (i.e. in `cases` environment)
-- or dummy boxes for null delimiters. 

local function real_box(list)
    for n, id, sb in traverse(list) do
        if id == glyph_id then
            return true
        elseif (id == hlist_id or id == vlist_id) and not ignored_subtypes[sb] then
            if real_box(n.list) then
                return true
            end
        end
    end
    return false
end

-- If the first thing (that we care about) in a line is a glyph
-- we simply number it, in which case `true` is returned.
-- Otherwise the line is just a wrapper around boxes, so we collect every
-- `vlist` reachable through box-only chains into `found` (each as a
-- `{node, hoffset, voffset}` triple) and return that list, so the caller can
-- recurse into all of them. We can't assume there is only a single box to
-- recurse into: an `hbox` can hold several numbering targets, either stacked
-- vertically (pgf stacks the title and the body of a tcolorbox with
-- `\raise`/`\lower`) or side by side (the columns of `multicol`, `paracol`,
-- two-column layouts, ...).
--
-- For each box we record both its horizontal offset (for placing the number)
-- and its vertical offset `voffset`. Inside an `hlist` the `shift` field is a
-- {\em vertical} displacement, so we accumulate it on every hop. The caller uses
-- `voffset` to tell columns (boxes side by side, sharing a `voffset`) from a
-- vertical stack (boxes at distinct `voffset`s).
--
-- The one structure where the boxes of an `hbox` form a {\em single} line that
-- should be numbered once is the cells of a table row. Those are told apart
-- by the caller using the subtype of the line (`alignment`), not here.

local function real_line(list, parent, offset, voffset, found)
    found = found or {}
    voffset = voffset or 0
    for n, id, sb in traverse(list) do
        if id == glyph_id then

-- A leading glyph means this really is a text line, so number it as a whole.
-- A glyph encountered after we already collected boxes is just part of the
-- wrapper and is ignored.

            if #found == 0 then return true end
        elseif id == vlist_id and not ignored_subtypes[sb] and real_box(n.list) then
            found[#found+1] = {n, offset + rangedimensions(parent, list, n), voffset + n.shift}
        elseif id == hlist_id and not ignored_subtypes[sb] and real_box(n.list) then
            if real_line(n.list, n, offset + rangedimensions(parent, list, n),
                         voffset + n.shift, found) == true then
                return true
            end
        end
    end
    return found
end

-- This function finds the lines that needs to be numbered in a page.
-- It should be used right before shipout, but can be used on individual
-- boxes using the `processbox` key if needed (maybe special numbering order
-- is desired). 
-- When a line found, `lineno_callbacks` is called to number it.
local find_line
find_line = function(parent, list, column, offset, width)

    if get_attribute(parent, mark_attr) == lineno_marks.processed then return end
    set_attribute(parent, mark_attr, lineno_marks.processed)
    
-- We need to keep track of the parent id to know if the `.shift`
-- field represent horizontal or vertical displacement.

    local parent_is_vlist = parent.id == vlist_id
    for n, id, sb in traverse(list) do
    
-- Lines are `hlist`s, so if a node is not one we dig deeper,
-- while calculating the offset and the width.
-- If a box marked with the `anchor` key is found
-- then the offset is reset and the width is updated.

        if id ~= hlist_id then
            if not n.list then goto continue end
            local new_offset, new_width = offset, width
            if get_attribute(n, mark_attr) == lineno_marks.anchor then
                new_offset, new_width = 0, n.width
            elseif parent_is_vlist then
                new_offset = new_offset + n.shift
            end
            find_line(n, n.list, column, new_offset, new_width)
            goto continue
        end
        
-- A line type is determined by the attribute of its last node so that line types can be switched
-- from within the line (but maybe this should be configurable).
-- The flag is a bitset that determines whether to number or recurse further.

        local line_attr = n.head and get_attribute(tail(n.head), type_attr)
        local line_type = line_attr and lineno_types[line_attr] and lineno_types[line_attr][column]
        local flag
        if sb == align_sub and get_attribute(n,mark_attr) == lineno_marks.display then
            flag = line_type and line_type[displayalign_sub]
        else
            flag = line_type and line_type[sb]
        end

-- If a line does not have any attribute we don't number it, but we do recurse further.

        local should_number = flag and (flag & LINENO_NUMBER) ~= 0 or false
        local should_recurse = flag and (flag & LINENO_RECURSE) ~= 0 or true 
        if not (should_number or should_recurse) then
            goto continue
        end
        
-- This is the case where a line should be numbered only once.
-- Maybe someone would like to number alignment once, regardless
-- of the fact the first column contains cells with paragraphs.        

        if should_number and not should_recurse then
            if real_box(n.list) then
                local new_offset = parent_is_vlist and (offset + n.shift) or offset
                lineno_callbacks(n, line_type, new_offset, width, parent.dir)
            end
            goto continue
        end
        
-- If `real_line` did not return `true`, the line is a wrapper around boxes,
-- so we need to find lines inside of each collected box as well. As before
-- offset and width might need to be updated.
--
-- The boxes of a table row are the exception: they are the cells of a single
-- line and must be numbered once, so for an alignment we only keep the first
-- box (this is what `real_line` used to return before it learned to collect
-- several boxes).
--
-- Otherwise we look at the vertical offsets of the collected boxes: boxes
-- sharing a `voffset` sit side by side and are columns, so each is processed
-- relative to its own left edge and width, and gets a column number assigned
-- in the order they appear (\LuaTeX/ keeps the columns in logical order).
-- Boxes alone at their `voffset` are a vertical stack (or just a single box)
-- and keep the inherited offset, width and column.

        local found = real_line(n.head, n, offset)
        if found ~= true then
            local last = sb == align_sub and math.min(#found, 1) or #found

-- A `voffset` shared by more than one box marks a row of columns; we count
-- the boxes per `voffset` to recognize them and number them within the row.

            local per_voffset = {}
            for i = 1, last do
                local v = found[i][3]
                per_voffset[v] = (per_voffset[v] or 0) + 1
            end
            local col_number = {}
            for i = 1, last do
                local box = found[i]
                local m, v = box[1], box[3]
                local new_offset, new_width, new_col = box[2], width
                if per_voffset[v] > 1 then
                    col_number[v] = (col_number[v] or 0) + 1
                    new_col = col_number[v]
                    new_offset, new_width = 0, m.width
                elseif get_attribute(m, mark_attr) == lineno_marks.anchor then
                    new_offset, new_width = 0, m.width
                elseif parent_is_vlist then
                    new_offset = new_offset + n.shift
                end
                find_line(m, m.head, new_col or column, new_offset, new_width)
            end
            goto continue
        end

        if not should_number then goto continue end

-- A line is found! Update the offset and number it.

        local new_offset = parent_is_vlist and (offset + n.shift) or offset
        lineno_callbacks(n, line_type, new_offset, width, parent.dir)

        ::continue::
    end
end

if not plain then
    luatexbase.add_to_callback('pre_shipout_filter', function(box)
        find_line(box, box.list, 1, 0, box.width)
        return true
    end, 'lualineno.shipout')
end

-- Check if inside display for display alignment.

if luatexbase then 
    luatexbase.add_to_callback("buildpage_filter", function(info)
        if info == "before_display" then 
            setattribute(mark_attr,lineno_marks.display)
        elseif info == "after_display" then
            setattribute(mark_attr,unset_attr)
        end
    end, "lualineno.indisplay")
end

-- \seccc Anchoring numbers to a box^^M

local function mark_last_vlist(n)
    local current = n
    while current do
        if current.id == vlist_id then
            set_attribute(current, mark_attr, lineno_marks.anchor)
            return true
        elseif current.id == hlist_id then
            if mark_last_vlist(tail(current.list)) then return true end
        end
        current = current.prev
    end
    return false
end

-- \seccc labels^^M

-- A label is attached by inserting a `user_defined` whatsit carrying the
-- label tokens as a lua table at the point of the `\lualineno{label=...}`
-- call.

local make_label
local function find_label(line)
    local head = line.list
    for n, id, sb in traverse(head) do
        if n.list then
            find_label(n)
        elseif id == whatsit_id and sb == user_defined_sub and n.user_id == label_whatsit then
            make_label(n.value, head, n)
        end
    end
end

-- The carrier stores the label tokens directly in the whatsit `value` field
-- (a lua table, type `l`). On its first use the function registers
-- `find_label` with the numbering callback and then redefines itself to the
-- bare inserter.

local add_label
add_label = function(tokens)
    if optex then
        luatexbase.add_to_callback('lualineno.pre_numbering', find_label, 'lualineno.labels')
    elseif latex then
        luatexbase.add_to_callback('lualineno.post_numbering', find_label, 'lualineno.labels')
    end
    add_label = function(toks)
        local w = node_copy(base_whatsit)
        w.value = toks
        w.attr = node.current_attr()
        node_write(w)
    end
    return add_label(tokens)
end

-- \seccc User Interface^^M
-- This section describes the definition of
-- the one macro exposed to the end user.
-- It is based on the luakeyval module.

local defaults = {
    toks = { },
    left = { },
    right = { },
    box = {number = true, recurse = true},
    alignment = {number = true, recurse = true},
    displayalignment = {number = true, recurse = true},
    equation = {number = true, recurse = true},
    line = {number = true, recurse = true},
    offset = true,
    column = 1,
}

local inner_keys = {
    number = {scanner = scan_bool, default = true},
    recurse = {scanner = scan_bool, default = true}
}

local defaults_keys = {
    toks = {scanner = scan_toks},
    left = {scanner = scan_toks},
    right = {scanner = scan_toks},
    box = {scanner = process_keys, args = {inner_keys, messages}},
    alignment = {scanner = process_keys, args = {inner_keys, messages}},
    displayalignment = {scanner = process_keys, args = {inner_keys, messages}},
    equation = {scanner = process_keys, args = {inner_keys, messages}},
    line = {scanner = process_keys, args = {inner_keys, messages}},
    offset = {scanner = scan_bool},
    column = {scanner = scan_int}
}

local function set_defaults()
    local vals = process_keys(defaults_keys,messages)
    for k,v in pairs(vals) do
        defaults[k] = v
    end
end

local define_keys = { }
for k,v in pairs(defaults_keys) do
    define_keys[k] = v
end
define_keys.name = {scanner = scan_string}

local function define_lineno()
    local vals = process_keys(define_keys, messages)
    local name = vals['name']
    if not name then 
        texerror("lualineno: Missing name when defining a lineno")
        return
    end
    
    local col = vals['column'] or defaults.column
    lineno_attrs[name] = lineno_attrs[name] or #lineno_types + 1
    local i = lineno_attrs[name]
    lineno_types[i] = lineno_types[i] or {}
    lineno_types[i][col] = lineno_types[i][col] or {}
    
    local c = lineno_types[i][col]
    
    local function store_type(key, subtype_id)
        local setting = vals[key] or defaults[key]
        local flags = 0
        if setting.number then flags = flags | LINENO_NUMBER end
        if setting.recurse then flags = flags | LINENO_RECURSE end
        c[subtype_id] = flags
    end
    
    store_type('box', box_sub)
    store_type('alignment', align_sub)
    store_type('equation', eq_sub)
    store_type('line', line_sub)
    store_type('displayalignment', displayalign_sub)
    
    c.toks = vals.toks or defaults.toks
    c.left = vals.left or defaults.left
    c.right = vals.right or defaults.right
    if vals.offset ~= nil then
        c.offset = vals.offset
    else
        c.offset = defaults.offset
    end
end

local lualineno_keys = {
    set = {scanner = scan_string},
    unset = { default = true },
    define = {scanner = function() return true end, func = define_lineno},
    defaults = {scanner = function() return true end, func = set_defaults},
    anchor = { default = true },
    label = {scanner = scan_toks, args = {false, true}},
    typeattr = {scanner = scan_int},
    markattr = {scanner = scan_int},
    whatsitid = {scanner = scan_int},
    processbox = {scanner = scan_int},
}

local function lualineno()
    local saved_endlinechar = tex.endlinechar
    tex.endlinechar = 32
    local vals = process_keys(lualineno_keys,messages)
    tex.endlinechar = saved_endlinechar
    if vals.set then
        local attr = lineno_attrs[vals.set]
        if attr then
          setattribute(type_attr, attr)
        else
          texerror("lualineno: type '" .. vals.set .. "' undefined")
        end
    end
    if vals.unset then
        setattribute(type_attr, unset_attr)
    end
    if vals.anchor then
        for i=texnest.ptr,0,-1 do 
            if mark_last_vlist(texnest[i].tail) then return end
        end
    end
    if vals.label then
        add_label(vals.label)
    end
    type_attr = vals.typeattr or type_attr
    mark_attr = vals.markattr or mark_attr
    label_whatsit = vals.whatsitid or label_whatsit
    if vals.processbox then
        local box = tex.box[vals.processbox]
        find_line(box, box.head, 1, 0, box.width)
    end
end

do
  if token.is_defined('lualineno') then
      texio.write_nl('log', "lualineno: redefining \\lualineno")
  end
  local function_table = lua.get_functions_table()
  local luafnalloc = luatexbase and luatexbase.new_luafunction 
    and luatexbase.new_luafunction('lualineno') or #function_table + 1
  token.set_lua('lualineno', luafnalloc, 'protected')
  function_table[luafnalloc] = lualineno
end

-- \seccc Format Specific Code^^M

if format == 'optex' then

-- To be able to use \OpTeX/'s color mechanism in line numbers the colorizing
-- needs to happen after line numbers are added, so we remove and insert
-- back again the colorizing function from the `pre_shipout_filter` callback.

    local colorize = callback.remove_from_callback('pre_shipout_filter', '_colors')
    callback.add_to_callback('pre_shipout_filter', colorize, '_colors')

-- \OpTeX/ only needs to run `\label[toks]` before a destination to label it.

    local lbracket = new_tok(string.byte('['), token.command_id'other_char')
    local rbracket = new_tok(string.byte(']'), token.command_id'other_char')
    local label_tok = create('_label')
    make_label = function(label)
        runtoks(function()
            put_next({rbracket})
            put_next(label)
            put_next({label_tok,lbracket})
        end)
    end
    
elseif latex then

-- If the luacolor package is loaded,
-- colorizing must happen after line numbers
-- are added to be able to color them.

    luatexbase.declare_callback_rule('pre_shipout_filter', 
         'lualineno.shipout', 'before', 'luacolor.process')

-- Since \LaTeX/ isn't really shipping out the page box, but a box
-- containing the `\topmargin` and the page box which is shifted with
-- `\moveright`, so it adds an undesired offset in `find_line`, so
-- we mark the page box as `anchor`.

    local attr_num = luatexbase.attributes['lualineno_mark']
    local replace = string.format([[\moveright \@themargin \vbox attr %d = %d]],
                                   attr_num,lineno_marks.anchor)
    local find = [[\moveright \@themargin \vbox]]
    local patch, num_subs = token.get_macro("@outputpage"):gsub(find, replace)

-- Log the success or failure of the patch.

    if num_subs > 0 then
        token.set_macro("@outputpage", patch)
    else
        texio.write_nl('log', "lualineno: failed to patch \\@outputpage")
    end

-- \LaTeX/'s `\label`'s creates a whatsit node (`\write`), so we temporarily box the label
-- to fetch this node, and add it to the list.

    local label_tok = create('label')
    make_label = function(label, list, n)
        runtoks(function()
            put_next({rbrace,rbrace})
            put_next(label)
            put_next({hbox, lbrace, label_tok,lbrace})
            local label_node = scan_list()
            list = insert_after(list,n,node_copy(label_node.head))
            node_flush(label_node)
        end)
    end
    
end