spb/metrika Public

1//2//  AST.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910/// A variable specification appearing in a varlist. Factor notation is11/// preserved verbatim; expansion into indicator columns happens at plan12/// time (CLAUDE.md §4), never at parse time.13public indirect enum ZQVarSpec: Equatable, Sendable {14    /// Plain variable reference: `price`15    case simple(String)16    /// Factor-operator variable: `i.region`, `c.age`, `ib2.group`17    case factor(op: String, name: String)18    /// Interaction: `c.age#c.age`, `i.region#c.price`19    case interaction([ZQVarSpec], full: Bool)2021    /// The underlying dataset column names this spec references.22    public var referencedNames: [String] {23        switch self {24        case .simple(let n): return [n]25        case .factor(_, let n): return [n]26        case .interaction(let parts, _): return parts.flatMap(\.referencedNames)27        }28    }29}3031/// Expression AST used for `if` qualifiers, `gen`/`replace` right-hand32/// sides, and option arguments.33public indirect enum ZQExpression: Equatable, Sendable {34    case number(Double)35    case string(String)36    case variable(String)37    /// Stata-style missing-value literal `.`38    case missing39    case unary(op: String, ZQExpression)40    case binary(op: String, ZQExpression, ZQExpression)41    case call(name: String, arguments: [ZQExpression])42}4344/// Observation range qualifier: `in 5/200`. Negative values count from the45/// end (`in -10/-1`), 1-based like Stata.46public struct ZQRange: Equatable, Sendable {47    public var lower: Int48    public var upper: Int4950    public init(lower: Int, upper: Int) {51        self.lower = lower52        self.upper = upper53    }54}5556/// Weight qualifier: `[aweight = expr]`.57public struct ZQWeight: Equatable, Sendable {58    public enum Kind: String, Equatable, Sendable {59        case aweight, fweight, pweight, iweight60    }6162    public var kind: Kind63    public var expression: ZQExpression6465    public init(kind: Kind, expression: ZQExpression) {66        self.kind = kind67        self.expression = expression68    }69}7071/// A trailing option: bare flag (`robust`, `detail`) or parenthesized72/// (`cluster(id)`, `reps(100000)`, `by(region year)`).73public struct ZQOption: Equatable, Sendable {74    public var name: String75    /// Raw argument tokens re-rendered as strings; interpretation is76    /// command-specific and happens at plan/execution time.77    public var arguments: [String]7879    public init(name: String, arguments: [String] = []) {80        self.name = name81        self.arguments = arguments82    }8384    public var firstArgument: String? { arguments.first }85}8687/// Instrumental-variables group in a varlist:88/// `ivregress 2sls y exog (endog1 endog2 = instr1 instr2)`.89public struct ZQIVSpec: Equatable, Sendable {90    public var endogenous: [String]91    public var instruments: [String]9293    public init(endogenous: [String], instruments: [String]) {94        self.endogenous = endogenous95        self.instruments = instruments96    }97}9899/// Assignment payload for `gen`/`replace`-style commands:100/// `gen log_rev = ln(revenue)`.101public struct ZQAssignment: Equatable, Sendable {102    public var target: String103    public var expression: ZQExpression104105    public init(target: String, expression: ZQExpression) {106        self.target = target107        self.expression = expression108    }109}110111/// A fully parsed ZQL command:112/// `command [varlist] [if expr] [in range] [weight] [, options]`,113/// optionally wrapped by a prefix command (`bootstrap, reps(1000): reg …`).114public struct ZQCommand: Equatable, Sendable {115    /// Canonical verb (abbreviations already resolved: `reg` → `regress`).116    public var verb: String117    /// Sub-verb for compound commands (`graph scatter …` → "scatter").118    public var subverb: String?119    public var varlist: [ZQVarSpec]120    /// Instrumented group for IV estimators, parsed from121    /// `(endog… = instruments…)` inside the varlist.122    public var ivSpec: ZQIVSpec?123    /// `gen`/`replace` assignment, mutually exclusive with a plain varlist.124    public var assignment: ZQAssignment?125    /// Raw argument such as a file path (`use sales.parquet`).126    public var argument: String?127    public var condition: ZQExpression?128    public var range: ZQRange?129    public var weight: ZQWeight?130    public var options: [ZQOption]131    /// Inner command of a prefix construct: for132    /// `bootstrap, reps(1000): reg y x`, the outer command is `bootstrap`133    /// and `body` holds the parsed `regress` command.134    public var body: ZQCommand? {135        get { bodyBox?.value }136        set { bodyBox = newValue.map(Box.init) }137    }138139    private var bodyBox: Box<ZQCommand>?140141    public init(142        verb: String,143        subverb: String? = nil,144        varlist: [ZQVarSpec] = [],145        ivSpec: ZQIVSpec? = nil,146        assignment: ZQAssignment? = nil,147        argument: String? = nil,148        condition: ZQExpression? = nil,149        range: ZQRange? = nil,150        weight: ZQWeight? = nil,151        options: [ZQOption] = [],152        body: ZQCommand? = nil153    ) {154        self.verb = verb155        self.subverb = subverb156        self.varlist = varlist157        self.ivSpec = ivSpec158        self.assignment = assignment159        self.argument = argument160        self.condition = condition161        self.range = range162        self.weight = weight163        self.options = options164        self.bodyBox = body.map(Box.init)165    }166167    public func option(_ name: String) -> ZQOption? {168        options.first { $0.name == name }169    }170171    public func hasOption(_ name: String) -> Bool {172        option(name) != nil173    }174}175176/// Reference box that lets `ZQCommand` nest itself as a prefix body.177private final class Box<T>: @unchecked Sendable {178    let value: T179    init(_ value: T) { self.value = value }180}181182extension Box: Equatable where T: Equatable {183    static func == (lhs: Box<T>, rhs: Box<T>) -> Bool { lhs.value == rhs.value }184}185