pdf-odin-llm/parser.odin

package pdf

import "core:mem"

// Token stream to Object (ISO 32000-2, 7.3).
//
// The parser is written so that no input can make it fail outright: an
// unterminated array yields the elements it did collect, a dictionary key that
// is not a name is dropped, a stream whose /Length lies is measured by finding
// its `endstream`. Each of those records a `Repair` on the Document and
// carries on, because a file with one broken object usually still renders. An
// `Error` comes back only when there is genuinely no object to hand over.
//
// A nil Document is supported throughout: inline image dictionaries in a
// content stream are parsed with no file behind them.

Parser :: struct {
    lex:       Lexer,
    doc:       ^Document,
    // Everything the parser builds is allocated here and never freed
    // individually, so this wants to be an arena -- `doc.allocator` for
    // anything that belongs to the Document.
    allocator: mem.Allocator,
    depth:     int,
    // The indirect object being parsed, which is what its strings and stream
    // are decrypted with. The zero Ref means "do not decrypt": that is the
    // trailer, an inline image, and every object that came out of an object
    // stream (those were decrypted with their container, ISO 32000-2, 7.5.7).
    obj_ref:   Ref,

    // Token lookahead. Two are needed and two are enough: telling "12 0 R"
    // from "12 0 obj" and from a plain "12" is the only place in the grammar
    // that needs to see past the token in hand.
    lookahead:   [2]Token,
    n_lookahead: int,
}

// parser_make prepares a parser over `data`.
//
// `data` is normally the whole file with `lex.pos` moved to the object, not a
// sub-slice: token offsets are then file offsets, which is what repairs
// record, and a stream payload can be a sub-slice of `doc.data` instead of a
// copy.
parser_make :: proc(doc: ^Document, data: []byte, allocator := context.allocator) -> Parser {
    p := Parser{
        lex       = lexer_make(data),
        doc       = doc,
        allocator = allocator,
    }
    if p.allocator.procedure == nil {
        p.allocator = context.allocator
    }
    return p
}

// ---------------------------------------------------------------------------
// Token queue
// ---------------------------------------------------------------------------

@(private = "file")
parser_next :: proc(p: ^Parser) -> Token {
    if p.n_lookahead > 0 {
        tok := p.lookahead[0]
        p.lookahead[0] = p.lookahead[1]
        p.n_lookahead -= 1
        return tok
    }
    return lexer_next(&p.lex, p.allocator)
}

// parser_peek looks `n` tokens ahead (0 or 1) without consuming anything.
@(private = "file")
parser_peek :: proc(p: ^Parser, n: int) -> Token {
    idx := clamp(n, 0, len(p.lookahead) - 1)
    for p.n_lookahead <= idx {
        p.lookahead[p.n_lookahead] = lexer_next(&p.lex, p.allocator)
        p.n_lookahead += 1
    }
    return p.lookahead[idx]
}

// token_ends_container names the keywords that can only appear outside of an
// array or dictionary. Meeting one means a closing bracket was lost, and
// swallowing it would take the rest of the file with it.
@(private = "file")
token_ends_container :: proc(t: Token) -> bool {
    if t.kind != .Keyword { return false }
    switch t.text {
    case "endobj", "stream", "endstream", "obj", "xref", "trailer", "startxref":
        return true
    }
    return false
}

// ---------------------------------------------------------------------------
// Objects
// ---------------------------------------------------------------------------

// parse_object reads one object. Indirect references are returned as `Ref`;
// nothing is resolved here.
parse_object :: proc(p: ^Parser) -> (Object, Error) {
    tok := parser_next(p)

    switch tok.kind {
    case .EOF:
        return nil, .Truncated

    case .Integer:
        // "12 0 R" is a reference; "12 0 obj" is a header that ran into us,
        // and "12 0" is two integers. Only the full three-token shape counts.
        t1 := parser_peek(p, 0)
        if t1.kind == .Integer && t1.int_value >= 0 && t1.int_value <= 65535 {
            t2 := parser_peek(p, 1)
            if t2.kind == .Keyword && t2.text == "R" {
                p.n_lookahead = 0
                return make_ref(tok.int_value, t1.int_value), .None
            }
        }
        return tok.int_value, .None

    case .Real:
        return tok.real_value, .None

    case .String:
        if tok.truncated {
            add_repair(p.doc, .Unterminated_String, i64(tok.start), p.obj_ref)
        }
        return String(parser_decrypt_string(p, tok.bytes)), .None

    case .Name:
        return Name(tok.text), .None

    case .Array_Open:
        return parse_array(p)

    case .Dict_Open:
        return parse_dict_or_stream(p)

    case .Keyword:
        switch tok.text {
        case "true":  return true, .None
        case "false": return false, .None
        case "null":  return nil, .None
        }
        return nil, .Unexpected_Token

    case .Array_Close, .Dict_Close, .Brace_Open, .Brace_Close:
        return nil, .Unexpected_Token

    case .Error:
        return nil, .Syntax_Error
    }

    return nil, .Syntax_Error
}

// make_ref validates an indirect reference. Object number 0 is the free-list
// head and can never be referenced, and a number past the format's limit
// cannot be in any xref table; both read as null, which is what a reference to
// a non-existent object means anyway (ISO 32000-2, 7.3.10).
@(private = "file")
make_ref :: proc(num, gen: i64) -> Object {
    if num < 1 || num > MAX_OBJECT_NUMBER || gen < 0 || gen > 65535 {
        return nil
    }
    return Ref{num = u32(num), gen = u16(gen)}
}

@(private = "file")
parser_decrypt_string :: proc(p: ^Parser, raw: []byte) -> []byte {
    if p.doc == nil || p.doc.encrypt == nil || p.obj_ref.num == 0 {
        return raw
    }
    return decrypt_string(p.doc, p.obj_ref, raw)
}

// ---------------------------------------------------------------------------
// Arrays
// ---------------------------------------------------------------------------

@(private = "file")
parse_array :: proc(p: ^Parser) -> (Object, Error) {
    if p.depth >= MAX_DEPTH {
        return nil, .Recursion_Limit
    }
    p.depth += 1
    defer p.depth -= 1

    a := new_array(p.allocator)
    // One token consumes at least one byte, so the remaining input bounds the
    // number of iterations. The counter is belt and braces against a token
    // rule that stops advancing.
    fuel := len(p.lex.data) - p.lex.pos + 8

    for fuel > 0 {
        fuel -= 1
        tok := parser_peek(p, 0)
        switch {
        case tok.kind == .Array_Close:
            parser_next(p)
            return a, .None
        case tok.kind == .EOF:
            add_repair(p.doc, .Unterminated_Array, i64(tok.start), p.obj_ref)
            return a, .None
        case tok.kind == .Dict_Close, token_ends_container(tok):
            // A ']' was lost. Leave the token for the enclosing dictionary or
            // the object reader; taking it would lose that structure too.
            add_repair(p.doc, .Unterminated_Array, i64(tok.start), p.obj_ref)
            return a, .None
        }

        obj, err := parse_object(p)
        if err == .Recursion_Limit {
            return a, err
        }
        if err != .None {
            // Junk between elements: the token was consumed, so step over it.
            // A null *element* is kept -- array positions carry meaning.
            continue
        }
        array_push(a, obj)
    }
    return a, .None
}

// ---------------------------------------------------------------------------
// Dictionaries and streams
// ---------------------------------------------------------------------------

@(private = "file")
parse_dict_or_stream :: proc(p: ^Parser) -> (Object, Error) {
    d, err := parse_dict(p)
    if err != .None {
        // The nesting limit. Whatever was collected before it is not worth
        // handing back as if it were the whole object, and an Object holding a
        // nil ^Dict is not the null object.
        return nil, err
    }
    // A stream is a dictionary followed by `stream`, and only ever as the
    // whole body of an indirect object -- so only at the outermost level.
    if p.depth == 0 {
        tok := parser_peek(p, 0)
        if tok.kind == .Keyword && tok.text == "stream" {
            parser_next(p)
            return parse_stream(p, d, tok.end), .None
        }
    }
    return d, .None
}

@(private = "file")
parse_dict :: proc(p: ^Parser) -> (^Dict, Error) {
    if p.depth >= MAX_DEPTH {
        return nil, .Recursion_Limit
    }
    p.depth += 1
    defer p.depth -= 1

    d := new_dict(p.allocator)
    fuel := len(p.lex.data) - p.lex.pos + 8

    for fuel > 0 {
        fuel -= 1
        key_tok := parser_peek(p, 0)
        switch {
        case key_tok.kind == .Dict_Close:
            parser_next(p)
            return d, .None
        case key_tok.kind == .EOF:
            add_repair(p.doc, .Unterminated_Dict, i64(key_tok.start), p.obj_ref)
            return d, .None
        case token_ends_container(key_tok):
            add_repair(p.doc, .Unterminated_Dict, i64(key_tok.start), p.obj_ref)
            return d, .None
        }

        if key_tok.kind != .Name {
            parser_next(p)
            add_repair(p.doc, .Unterminated_Dict, i64(key_tok.start), p.obj_ref)
            // The value belongs to a key that cannot be recorded, so it goes
            // with it; left behind it would be read as the next key. A stray
            // bracket is junk on its own and takes nothing with it.
            #partial switch key_tok.kind {
            case .Array_Close, .Brace_Open, .Brace_Close, .Error:
            case:
                skip := parser_peek(p, 0)
                if skip.kind != .Dict_Close && skip.kind != .EOF && !token_ends_container(skip) {
                    _, _ = parse_object(p)
                }
            }
            continue
        }
        parser_next(p)

        // "/Key >>" -- a key whose value was lost. Closing the dictionary here
        // keeps the '>>' meaningful instead of consuming it as the value.
        vt := parser_peek(p, 0)
        if vt.kind == .Dict_Close || vt.kind == .EOF || token_ends_container(vt) {
            add_repair(p.doc, .Unterminated_Dict, i64(key_tok.start), p.obj_ref)
            if vt.kind == .Dict_Close {
                parser_next(p)
            }
            return d, .None
        }

        value, verr := parse_object(p)
        if verr == .Recursion_Limit {
            // The value is deeper than the library will follow, and its tokens
            // are still on the stream, so nothing after it can be trusted to
            // be a key. Fail the dictionary the way an array fails.
            return d, verr
        }
        if verr != .None {
            continue
        }
        // A key with a null value and an absent key mean the same thing
        // everywhere in the format, so nulls are dropped at parse time.
        if value != nil {
            dict_set(d, key_tok.text, value)
        }
    }
    return d, .None
}

// parse_stream reads the payload that follows a `stream` keyword ending at
// `kw_end`. The payload is a sub-slice of the input, never a copy -- except
// when it has to be decrypted.
@(private = "file")
parse_stream :: proc(p: ^Parser, d: ^Dict, kw_end: int) -> Object {
    data := p.lex.data

    // "stream" is followed by CRLF or LF, never by CR alone (ISO 32000-2,
    // 7.3.8.1). Producers that emit a bare CR, or a space before the EOL, are
    // common enough that both are tolerated.
    i := kw_end
    for i < len(data) && (data[i] == ' ' || data[i] == '\t') {
        i += 1
    }
    if i < len(data) && data[i] == '\r' { i += 1 }
    if i < len(data) && data[i] == '\n' { i += 1 }
    body := i

    length := i64(-1)
    if lv, ok := to_int(dict_get(p.doc, d, "Length")); ok {
        length = lv
    }

    end := -1
    after := -1
    // Trust /Length only if it fits in the file *and* `endstream` is where it
    // says it is. A length that is merely plausible is not evidence.
    if length >= 0 && length <= i64(len(data) - body) {
        probe := Lexer{data = data, pos = body + int(length)}
        if lexer_accept_keyword(&probe, "endstream") {
            end = body + int(length)
            after = probe.pos
        }
    }

    if end < 0 {
        needle := "endstream"
        idx := index_bytes(data[body:], transmute([]byte)needle)
        if idx < 0 {
            // No endstream at all: the file is truncated inside the stream.
            // Everything that is left is the best available payload.
            add_repair(p.doc, .Missing_Endstream, i64(body), p.obj_ref)
            end = len(data)
            after = len(data)
        } else {
            e := body + idx
            after = e + len(needle)
            // One EOL before `endstream` is the separator, not data. CRLF is
            // one marker, so both bytes come off.
            if e > body && data[e - 1] == '\n' { e -= 1 }
            if e > body && data[e - 1] == '\r' { e -= 1 }
            end = e
            add_repair(p.doc, .Bad_Stream_Length, i64(body), p.obj_ref)
        }
    }

    raw := data[body:end]
    // `raw` is a view of the file, so decrypt_stream has to return a new
    // buffer rather than writing through this one. An object that came out of
    // an object stream has the zero Ref and is already in the clear: its
    // container was decrypted as a whole (ISO 32000-2, 7.5.7).
    if p.doc != nil && p.doc.encrypt != nil && p.obj_ref.num != 0 {
        raw = decrypt_stream(p.doc, p.obj_ref, raw, d)
    }

    p.n_lookahead = 0
    p.lex.pos = clamp(after, 0, len(data))

    s := new_stream(d, raw, p.allocator)
    s.ref = p.obj_ref
    return s
}

// ---------------------------------------------------------------------------
// Indirect objects
// ---------------------------------------------------------------------------

// parse_indirect_object_at reads "num gen obj <object> endobj" at `offset`,
// an absolute index into `doc.data`.
//
// `got` is the reference the file actually has there, which is not always the
// one the cross-reference table promised; the caller compares them to detect a
// stale table. A header that is not exactly at `offset` is searched for
// nearby: incremental updaters that miscount by a byte or two are common
// enough that giving up on the first mismatch loses readable objects.
parse_indirect_object_at :: proc(doc: ^Document, offset: i64, want: Ref) -> (obj: Object, got: Ref, err: Error) {
    if doc == nil {
        return nil, {}, .Invalid_Argument
    }
    if offset < 0 || offset >= i64(len(doc.data)) {
        return nil, {}, .Object_Not_Found
    }

    start := int(offset)
    num, gen, body, ok := object_header_at(doc.data, start)
    if !ok {
        num, gen, body, ok = object_header_near(doc.data, start, want)
        // Offsets in a file with junk before "%PDF-" are sometimes measured
        // from the header instead of from byte zero.
        if !ok && doc.header_offset > 0 && offset + doc.header_offset < i64(len(doc.data)) {
            num, gen, body, ok = object_header_near(doc.data, int(offset + doc.header_offset), want)
        }
        if ok {
            add_repair(doc, .Xref_Entry_Mismatch, offset, want)
        }
    }
    if !ok {
        return nil, {}, .Object_Not_Found
    }

    got = Ref{num = num, gen = gen}

    p := parser_make(doc, doc.data, doc.allocator)
    p.lex.pos = body
    p.obj_ref = got

    // "12 0 obj endobj" is an empty object, which is null rather than an error.
    if lexer_at_keyword(&p.lex, "endobj") {
        return nil, got, .None
    }

    obj, err = parse_object(&p)
    if err != .None {
        return nil, got, err
    }
    return obj, got, .None
}

// object_header_at matches "num gen obj" exactly at `pos`.
@(private = "file")
object_header_at :: proc(data: []byte, pos: int) -> (num: u32, gen: u16, body: int, ok: bool) {
    if pos < 0 || pos >= len(data) {
        return 0, 0, 0, false
    }
    lex := Lexer{data = data, pos = pos}
    // A header is three tokens, two integers and a keyword, and this runs at
    // every byte of a 129-wide window for every object of a damaged file:
    // `lexer_next_structural` is what keeps a rejected offset from copying a
    // string body the probe is about to throw away.
    t1 := lexer_next_structural(&lex)
    if t1.kind != .Integer || t1.int_value < 1 || t1.int_value > MAX_OBJECT_NUMBER {
        return 0, 0, 0, false
    }
    t2 := lexer_next_structural(&lex)
    if t2.kind == .Keyword && t2.text == "obj" {
        // "12 obj": the generation number was dropped. Generation 0 is right
        // in every file this has been seen in.
        return u32(t1.int_value), 0, lex.pos, true
    }
    if t2.kind != .Integer || t2.int_value < 0 || t2.int_value > 65535 {
        return 0, 0, 0, false
    }
    t3 := lexer_next_structural(&lex)
    if t3.kind != .Keyword || t3.text != "obj" {
        return 0, 0, 0, false
    }
    return u32(t1.int_value), u16(t2.int_value), lex.pos, true
}

// OBJECT_HEADER_WINDOW is how far either side of a claimed offset a header is
// looked for. Wide enough for the off-by-a-few-bytes tables that broken
// incremental updates produce, narrow enough that it cannot wander into the
// next object.
@(private = "file")
OBJECT_HEADER_WINDOW :: 64

@(private = "file")
object_header_near :: proc(data: []byte, center: int, want: Ref) -> (num: u32, gen: u16, body: int, ok: bool) {
    lo := max(center - OBJECT_HEADER_WINDOW, 0)
    hi := min(center + OBJECT_HEADER_WINDOW, len(data) - 1)

    for i := lo; i <= hi; i += 1 {
        // Only start where a token can start, or "123 0 obj" also matches at
        // its second digit and yields object 23.
        if data[i] < '0' || data[i] > '9' { continue }
        if i > 0 && is_pdf_regular(data[i - 1]) { continue }

        n, g, b, hit := object_header_at(data, i)
        if !hit { continue }
        if want.num == 0 || n == want.num {
            return n, g, b, true
        }
        if !ok {
            // Remember the first header found, in case none matches: an object
            // with the wrong number is still something the caller can judge.
            num, gen, body, ok = n, g, b, true
        }
    }
    return
}