akribes_types/

ast.rs

//! AST shapes that travel over the wire.
//!
//! This is the SDK-facing slice of `akribes_core::ast`: just the types
//! that consumers need to interpret engine events ([`Span`], [`TypeRef`],
//! [`TypeField`], [`ActorHint`], [`FieldConstraint`] and the associated
//! sentinel constants). The full `akribes_core::ast` module also defines
//! `Stmt`, `Expr`, `Program`, and the rest of the language AST, which
//! stays in core because the parser/analyzer/compiler own those shapes.

use serde::{Deserialize, Serialize};
use std::fmt;

#[derive(Debug, Clone, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
pub struct Span {
    pub line: usize,
    pub col: usize,
    pub end_line: usize,
    pub end_col: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TypeField {
    pub name: String,
    pub ty: TypeRef,
    pub docs: Option<String>,
    pub span: Span,
    /// Field-level validation constraints declared via `matches /re/`,
    /// `at_least_items 3`, prose `"..."` lines, etc. Attached by the parser
    /// to the most recent field whose `span.col` matches the constraint's
    /// column (see `parser.rs` constraint attachment rules). `#[serde(default)]`
    /// so ASTs serialized before constraints existed still deserialize.
    #[serde(default)]
    pub constraints: Vec<FieldConstraint>,
}

/// A single field-level validation constraint attached to a `type` field.
/// Parsed from the Constraint Mini-Language (see
/// `docs/superpowers/specs/2026-04-18-epa-constraint-language-design.md`).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum FieldConstraint {
    /// A structured (Tier-1) constraint like `matches /^[A-Z0-9]+$/` or
    /// `at_least_items 3`. `phrase` is the canonical phrase name (e.g.
    /// `"matches"`) — the key the `ConstraintRegistry` uses to look up the
    /// handler. `args` is a handler-specific JSON payload (for `matches` this
    /// is `{"pattern": "<regex>"}`).
    Tier1 {
        phrase: String,
        args: serde_json::Value,
        span: Span,
    },
    /// An unrecognized Tier-2 prose rule, e.g.
    /// `"must be a valid ticker symbol"`. Rendered verbatim into the prompt
    /// in Tier-1 prose form; never enforced at runtime.
    ProseRule { text: String, span: Span },
    /// A `validate_with: <ident>` custom-validator hook. `name` is the
    /// validator's canonical identifier — resolved against the
    /// `validation::validator_registry::VALIDATORS` registry at
    /// analysis time (emits `AKRIBES-E-VALIDATE-WITH-UNKNOWN` on misses) and
    /// dispatched at task-end (failures surface as
    /// `AKRIBES-E-VALIDATE-WITH-FAIL` corrective retries).
    ValidateWith { name: String, span: Span },
}

impl FieldConstraint {
    pub fn span(&self) -> &Span {
        match self {
            FieldConstraint::Tier1 { span, .. } => span,
            FieldConstraint::ProseRule { span, .. } => span,
            FieldConstraint::ValidateWith { span, .. } => span,
        }
    }
}

/// Upper cap on discriminated-union arms. 8 is the tightest reliable
/// value across Anthropic tool-use + Gemini `responseSchema` under
/// preliminary testing. The analyzer raises `AKRIBES-E-UNION-009` when an
/// arm list exceeds this.
pub const MAX_UNION_ARMS: usize = 8;

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TypeRef {
    pub name: String,
    pub inner: Option<Box<TypeRef>>,
    /// Populated only for string-literal union types; in that case `name` is
    /// the sentinel `"choice"` and `choices` holds the variant strings in
    /// declaration order. `None` for every other type. Variant validation
    /// (non-empty, unique, ≥2) is the analyzer's job, not the parser's.
    pub choices: Option<Vec<String>>,
    /// Populated only for discriminated-union types (general `A | B | ...`
    /// including the binary `T | Unable` special case). When `Some`, `name`
    /// is the sentinel `"variant_union"` (mirroring `"choice"`), both
    /// `inner` and `choices` are `None`, and `variants` holds every arm in
    /// source order with length in `[2, MAX_UNION_ARMS]`. The analyzer
    /// enforces arm-record-only, ≤8, no duplicates, and
    /// return-position-only usage in v1.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub variants: Option<Vec<TypeRef>>,
    /// Source span covering the type reference's surface text. Populated
    /// by the parser for every `TypeRef` it constructs; synthesized refs
    /// (stdlib signatures, analyzer alias substitutions, test fixtures,
    /// …) default to `Span::default()`.
    ///
    /// For composite types (`list[T]`), the OUTER `TypeRef`'s span covers
    /// the whole `list[T]` text and the INNER `T`'s span covers just the
    /// inner identifier. For optional sentinels (`T?`), the inner `T`'s
    /// span is the pre-`?` text and the outer sentinel's span includes
    /// the `?`. For variant unions (`A | B | ...`) and choice strings
    /// (`"a" | "b"`), each arm carries its own span and the outer
    /// sentinel spans the whole union text.
    ///
    /// `#[serde(default)]` so AST payloads serialized before this field
    /// existed (older SDK clients, durable execution caches) keep
    /// deserializing cleanly.
    #[serde(default)]
    pub span: Span,
}

/// Sentinel `TypeRef.name` for a discriminated union (`A | B | ...`).
pub const VARIANT_UNION_SENTINEL: &str = "variant_union";

/// Sentinel `TypeRef.name` for an optional type (`T?`). The `inner` field
/// holds the wrapped `T`. `none` is assignable to any optional type;
/// `T?` is NOT assignable to `T` without an explicit `?? default` unwrap
/// or a pattern-match on `none`. (D2)
pub const OPTIONAL_SENTINEL: &str = "optional";

impl TypeRef {
    /// Build a primitive or named type reference (no generic inner, no
    /// choice variants). Use this in preference to a struct literal so the
    /// `choices` field stays consistently `None` at non-choice sites.
    ///
    /// The resulting `span` is `Span::default()` — appropriate for
    /// synthesized refs (stdlib signatures, analyzer substitutions, test
    /// fixtures). Parser sites that have a concrete source span should
    /// use [`TypeRef::primitive_with_span`] instead so hover /
    /// goto-definition land on the right text range.
    pub fn primitive(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            inner: None,
            choices: None,
            variants: None,
            span: Span::default(),
        }
    }

    /// Build a primitive or named type reference with an explicit source
    /// span. Parser-facing companion to [`TypeRef::primitive`].
    pub fn primitive_with_span(name: impl Into<String>, span: Span) -> Self {
        Self {
            name: name.into(),
            inner: None,
            choices: None,
            variants: None,
            span,
        }
    }

    /// Build an optional type `T?` wrapping `inner` (D2). Idempotent:
    /// applying this to a type that is already optional returns the same
    /// shape (no double-wrap), matching most languages' `T??` collapse.
    ///
    /// The outer sentinel's `span` is `Span::default()` — parser sites
    /// that have the postfix `?` location should use
    /// [`TypeRef::optional_with_span`] so the outer sentinel covers the
    /// whole `T?` text (the inner `T`'s span stays at its own range).
    pub fn optional(inner: TypeRef) -> Self {
        if inner.is_optional() {
            return inner;
        }
        Self {
            name: OPTIONAL_SENTINEL.to_string(),
            inner: Some(Box::new(inner)),
            choices: None,
            variants: None,
            span: Span::default(),
        }
    }

    /// Build an optional type `T?` with an explicit source span covering
    /// the whole `T?` text (inner `T`'s span is preserved). Parser-facing
    /// companion to [`TypeRef::optional`].
    pub fn optional_with_span(inner: TypeRef, span: Span) -> Self {
        if inner.is_optional() {
            return inner;
        }
        Self {
            name: OPTIONAL_SENTINEL.to_string(),
            inner: Some(Box::new(inner)),
            choices: None,
            variants: None,
            span,
        }
    }

    /// `true` iff this `TypeRef` is an `Optional[T]` sentinel.
    pub fn is_optional(&self) -> bool {
        self.name == OPTIONAL_SENTINEL && self.inner.is_some()
    }

    /// Borrow the wrapped `T` from an `Optional[T]`; `None` for non-optional.
    pub fn optional_inner(&self) -> Option<&TypeRef> {
        if self.is_optional() {
            self.inner.as_deref()
        } else {
            None
        }
    }

    /// Build a discriminated union from an ordered arm list. Grammar
    /// guarantees `arms.len() >= 2`; arm-count caps are analyzer-enforced
    /// (`AKRIBES-E-UNION-009`) so oversized unions reach the analyzer as
    /// parsed ASTs instead of panicking in the parser. The binary
    /// `T | Unable` case is just `variant_union(vec![T, Unable])`.
    pub fn variant_union(arms: Vec<TypeRef>) -> Self {
        debug_assert!(arms.len() >= 2, "variant union requires >= 2 arms");
        Self {
            name: VARIANT_UNION_SENTINEL.to_string(),
            inner: None,
            choices: None,
            variants: Some(arms),
            span: Span::default(),
        }
    }

    /// Build a discriminated union with an explicit outer span covering
    /// the whole `A | B | ...` text. Each arm's `TypeRef` keeps its own
    /// span (arms come from the parser already populated). Parser-facing
    /// companion to [`TypeRef::variant_union`].
    pub fn variant_union_with_span(arms: Vec<TypeRef>, span: Span) -> Self {
        debug_assert!(arms.len() >= 2, "variant union requires >= 2 arms");
        Self {
            name: VARIANT_UNION_SENTINEL.to_string(),
            inner: None,
            choices: None,
            variants: Some(arms),
            span,
        }
    }

    /// Build a binary union `success | Unable`. Kept as a named constructor
    /// because every #157 call site uses it; internally delegates to
    /// [`variant_union`] with `[success, Unable]` in canonical source
    /// order.
    pub fn union_with_unable(success: TypeRef) -> Self {
        Self::variant_union(vec![success, TypeRef::primitive("Unable")])
    }

    /// Return `true` iff this `TypeRef` is a discriminated-union sentinel
    /// (any arm count).
    pub fn is_variant_union(&self) -> bool {
        self.name == VARIANT_UNION_SENTINEL && self.variants.is_some()
    }

    /// Slice over the declared arms in source order (or `None` for
    /// non-union types).
    pub fn union_arms(&self) -> Option<&[TypeRef]> {
        self.variants.as_deref()
    }

    /// Return `true` iff this `TypeRef` is a binary union whose two arms
    /// are exactly one non-Unable record and one `Unable`. Used by every
    /// #157 call site that gates on "this is a T | Unable return type" —
    /// kept for backwards compatibility and cheap pattern-matching.
    pub fn is_union_with_unable(&self) -> bool {
        match self.variants.as_deref() {
            Some(arms) if arms.len() == 2 => {
                (arms[0].name == "Unable") ^ (arms[1].name == "Unable")
            }
            _ => false,
        }
    }

    /// Return the non-Unable branch of a binary `T | Unable`, or `None` if
    /// this is not exactly such a union. N-ary unions and unions without
    /// an `Unable` arm return `None` — callers that need the general arm
    /// list should use [`union_arms`].
    pub fn unwrap_union_success(&self) -> Option<&TypeRef> {
        match self.variants.as_deref() {
            Some(arms) if arms.len() == 2 => {
                if arms[0].name == "Unable" && arms[1].name != "Unable" {
                    Some(&arms[1])
                } else if arms[1].name == "Unable" && arms[0].name != "Unable" {
                    Some(&arms[0])
                } else {
                    None
                }
            }
            _ => None,
        }
    }

    /// Return the declared success arm of any discriminated union — the
    /// first arm in source order. Used for retry gating and
    /// `on <variant> default` type-checking. Returns `None` for non-union
    /// types.
    pub fn union_success_arm(&self) -> Option<&TypeRef> {
        self.variants.as_deref().and_then(|arms| arms.first())
    }

    /// Render a `TypeRef` as a source-level fragment for error messages,
    /// LSP labels, and user-facing diagnostics. Union types render as
    /// `A | B | ...`; choice types render as `"a" | "b" | ...`; generics
    /// render as `list[str]`; primitives render as their `name`.
    pub fn display(&self) -> String {
        // D2: `Optional[T]` renders as `T?` at the source level, mirroring
        // the postfix syntax authors typed.
        if let Some(inner) = self.optional_inner() {
            return format!("{}?", inner.display());
        }
        if let Some(arms) = &self.variants {
            return arms
                .iter()
                .map(|a| a.display())
                .collect::<Vec<_>>()
                .join(" | ");
        }
        if let Some(choices) = &self.choices {
            choices
                .iter()
                .map(|c| format!("\"{}\"", c))
                .collect::<Vec<_>>()
                .join(" | ")
        } else if let Some(inner) = &self.inner {
            format!("{}[{}]", self.name, inner.display())
        } else {
            self.name.clone()
        }
    }
}

impl fmt::Display for TypeRef {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(&self.display())
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ActorHint {
    Human,
    Any,
    Client(String),
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Span derives `Default` — required so existing `TypeRef { … }`
    /// struct literals across the workspace can keep building without
    /// every test fixture having to construct one by hand.
    #[test]
    fn span_default_is_all_zero() {
        let s = Span::default();
        assert_eq!(s.line, 0);
        assert_eq!(s.col, 0);
        assert_eq!(s.end_line, 0);
        assert_eq!(s.end_col, 0);
    }

    #[test]
    fn typeref_primitive_has_default_span() {
        let ty = TypeRef::primitive("str");
        assert_eq!(ty.span, Span::default());
    }

    #[test]
    fn typeref_primitive_with_span_carries_span() {
        let s = Span {
            line: 4,
            col: 7,
            end_line: 4,
            end_col: 10,
        };
        let ty = TypeRef::primitive_with_span("str", s.clone());
        assert_eq!(ty.span, s);
    }

    #[test]
    fn typeref_serde_roundtrip_preserves_span() {
        let s = Span {
            line: 2,
            col: 5,
            end_line: 2,
            end_col: 9,
        };
        let ty = TypeRef::primitive_with_span("str", s.clone());
        let json = serde_json::to_string(&ty).unwrap();
        let back: TypeRef = serde_json::from_str(&json).unwrap();
        assert_eq!(back.name, "str");
        assert_eq!(back.span, s);
    }

    /// The pre-`span` wire shape. Older SDK clients and durable execution
    /// caches serialized TypeRefs without a `span` field, so the new
    /// deserializer MUST accept those payloads and default-fill the
    /// span. This is the load-bearing compatibility guarantee.
    #[test]
    fn typeref_deserializes_legacy_payload_without_span_field() {
        let legacy = r#"{"name":"str","inner":null,"choices":null}"#;
        let ty: TypeRef = serde_json::from_str(legacy).unwrap();
        assert_eq!(ty.name, "str");
        assert_eq!(ty.span, Span::default());
    }

    #[test]
    fn typeref_deserializes_legacy_list_payload_without_span() {
        let legacy =
            r#"{"name":"list","inner":{"name":"str","inner":null,"choices":null},"choices":null}"#;
        let ty: TypeRef = serde_json::from_str(legacy).unwrap();
        assert_eq!(ty.name, "list");
        assert_eq!(ty.span, Span::default());
        let inner = ty.inner.as_deref().unwrap();
        assert_eq!(inner.name, "str");
        assert_eq!(inner.span, Span::default());
    }

    #[test]
    fn typeref_deserializes_legacy_choice_payload_without_span() {
        let legacy = r#"{"name":"choice","inner":null,"choices":["a","b"]}"#;
        let ty: TypeRef = serde_json::from_str(legacy).unwrap();
        assert_eq!(ty.name, "choice");
        assert_eq!(
            ty.choices.as_deref().unwrap(),
            &["a".to_string(), "b".to_string()]
        );
        assert_eq!(ty.span, Span::default());
    }

    /// Variant-union shapes already used `#[serde(default,
    /// skip_serializing_if = "Option::is_none")]` for `variants`, so a
    /// legacy payload without `variants` round-trips fine alongside a
    /// missing `span`. Lock that in.
    #[test]
    fn typeref_deserializes_legacy_payload_without_variants_or_span() {
        let legacy = r#"{"name":"str","inner":null,"choices":null}"#;
        let ty: TypeRef = serde_json::from_str(legacy).unwrap();
        assert_eq!(ty.name, "str");
        assert!(ty.variants.is_none());
        assert_eq!(ty.span, Span::default());
    }

    #[test]
    fn typeref_variant_union_roundtrip_preserves_arm_spans() {
        let arm_a = TypeRef::primitive_with_span(
            "A",
            Span {
                line: 1,
                col: 1,
                end_line: 1,
                end_col: 2,
            },
        );
        let arm_b = TypeRef::primitive_with_span(
            "B",
            Span {
                line: 1,
                col: 5,
                end_line: 1,
                end_col: 6,
            },
        );
        let outer = TypeRef::variant_union_with_span(
            vec![arm_a, arm_b],
            Span {
                line: 1,
                col: 1,
                end_line: 1,
                end_col: 6,
            },
        );
        let json = serde_json::to_string(&outer).unwrap();
        let back: TypeRef = serde_json::from_str(&json).unwrap();
        assert!(back.is_variant_union());
        let arms = back.union_arms().unwrap();
        assert_eq!(arms[0].name, "A");
        assert_eq!(arms[0].span.col, 1);
        assert_eq!(arms[1].name, "B");
        assert_eq!(arms[1].span.col, 5);
        assert_eq!(back.span.end_col, 6);
    }

    #[test]
    fn typeref_optional_with_span_returns_inner_if_already_optional() {
        // Idempotent collapse with a different outer span: the second
        // wrap is a no-op and returns the first-level optional
        // unchanged.
        let inner = TypeRef::primitive_with_span(
            "int",
            Span {
                line: 1,
                col: 1,
                end_line: 1,
                end_col: 4,
            },
        );
        let first = TypeRef::optional_with_span(
            inner,
            Span {
                line: 1,
                col: 1,
                end_line: 1,
                end_col: 5,
            },
        );
        let twice = TypeRef::optional_with_span(
            first.clone(),
            Span {
                line: 9,
                col: 9,
                end_line: 9,
                end_col: 9,
            },
        );
        assert_eq!(twice.span, first.span);
    }
}