use indexmap::IndexMap;
use petgraph::graph::{DiGraph, NodeIndex};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::diag::{Diagnostic, Span};
use crate::types::{ActiveType, Constraint, Unit};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct NodeId(pub u32);
#[derive(Debug, Clone)]
pub struct Module {
pub arena: Vec<Node>,
pub bindings: HashMap<String, NodeId>,
pub graph: DiGraph<NodeId, ()>,
pub node_to_graph: HashMap<NodeId, NodeIndex>,
pub root: Option<NodeId>,
pub schemas: HashMap<String, NodeId>,
pub diagnostics: Vec<Diagnostic>,
pub source_name: String,
pub node_spans: HashMap<NodeId, Span>,
pub node_files: HashMap<NodeId, String>,
pub healing: Vec<Diagnostic>,
pub trivia: Vec<TriviaComment>,
pub source_style: Style,
pub includes: Vec<IncludeStmt>,
pub profiles: IndexMap<String, NodeId>,
pub patches: Vec<NodeId>,
pub schema: Option<NodeId>,
pub annotations: HashMap<NodeId, Vec<Annotation>>,
pub depends_on: HashMap<NodeId, DependsOn>,
}
#[derive(Debug, Clone)]
pub struct Annotation {
pub kind: AnnotationKind,
pub path: String,
pub span: Span,
}
#[derive(Debug, Clone)]
pub enum AnnotationKind {
Intent(String),
Constraint { source: String, expr: NodeId },
AiImmutable(bool),
}
#[derive(Debug, Clone)]
pub struct DependsOn {
pub paths: Vec<String>,
pub span: Span,
}
#[derive(Debug, Clone)]
pub struct IncludeStmt {
pub path: String,
pub alias: Option<String>,
pub private: bool,
pub scope_path: Vec<String>,
pub line: u32,
}
#[derive(Debug, Clone)]
pub struct SchemaField {
pub typ: ActiveType,
pub min: Option<NodeId>,
pub max: Option<NodeId>,
pub min_len: Option<NodeId>,
pub max_len: Option<NodeId>,
pub pattern: Option<NodeId>,
pub default: Option<NodeId>,
pub deprecated: Option<String>,
pub span: Option<crate::diag::Span>,
}
impl SchemaField {
pub fn bare(typ: ActiveType, span: Option<crate::diag::Span>) -> Self {
Self {
typ,
min: None,
max: None,
min_len: None,
max_len: None,
pattern: None,
default: None,
deprecated: None,
span,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Style {
Brace,
Indent,
Preserve,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Attachment {
Leading(String),
Trailing(String),
TrailingBlock(String),
TrailingDocument,
}
#[derive(Debug, Clone)]
pub struct TriviaComment {
pub text: String,
pub span: Span,
pub attachment: Attachment,
pub file: Option<String>,
}
impl Module {
pub fn new(source_name: impl Into<String>) -> Self {
Self {
arena: Vec::new(),
bindings: HashMap::new(),
graph: DiGraph::new(),
node_to_graph: HashMap::new(),
root: None,
schemas: HashMap::new(),
diagnostics: Vec::new(),
source_name: source_name.into(),
node_spans: HashMap::new(),
node_files: HashMap::new(),
healing: Vec::new(),
trivia: Vec::new(),
source_style: Style::Brace,
includes: Vec::new(),
profiles: IndexMap::new(),
patches: Vec::new(),
schema: None,
annotations: HashMap::new(),
depends_on: HashMap::new(),
}
}
pub fn add_node(&mut self, node: Node) -> NodeId {
let id = NodeId(self.arena.len() as u32);
self.arena.push(node);
let idx = self.graph.add_node(id);
self.node_to_graph.insert(id, idx);
id
}
pub fn add_edge(&mut self, from: NodeId, to: NodeId) {
if let (Some(&f), Some(&t)) = (self.node_to_graph.get(&from), self.node_to_graph.get(&to)) {
self.graph.add_edge(f, t, ());
}
}
pub fn get(&self, id: NodeId) -> Option<&Node> {
self.arena.get(id.0 as usize)
}
pub fn get_mut(&mut self, id: NodeId) -> Option<&mut Node> {
self.arena.get_mut(id.0 as usize)
}
pub fn record_span(&mut self, id: NodeId, span: Span) {
self.node_spans.insert(id, span);
}
pub fn record_file(&mut self, id: NodeId, file: String) {
self.node_files.insert(id, file);
}
pub fn value_paths(&self) -> HashMap<NodeId, String> {
let mut paths = HashMap::new();
for (name, id) in &self.bindings {
paths.insert(*id, name.clone());
}
if let Some(root) = self.root {
let mut stack = vec![(root, String::new())];
while let Some((id, path)) = stack.pop() {
if let Some(Node::Table { fields }) = self.get(id) {
for (key, child) in fields {
let child_path = if path.is_empty() {
key.clone()
} else {
format!("{path}.{key}")
};
paths.entry(*child).or_insert_with(|| child_path.clone());
stack.push((*child, child_path));
}
}
}
}
paths
}
pub fn resolve_value_path(&self, path: &str) -> Option<NodeId> {
if let Some(root) = self.root {
let mut id = root;
let mut found = true;
for segment in path.split('.') {
match self.get(id) {
Some(Node::Table { fields }) => match fields.get(segment) {
Some(child) => id = *child,
None => {
found = false;
break;
}
},
_ => {
found = false;
break;
}
}
}
if found {
return Some(id);
}
}
if !path.contains('.') {
return self.bindings.get(path).copied();
}
None
}
}
#[derive(Debug, Clone)]
pub enum Node {
Literal(Value),
Env {
name: String,
default: Option<NodeId>,
},
Secret {
locator: NodeId,
},
Ref(NodeId),
Ident(String), Interpolate {
parts: Vec<InterpPart>,
},
Binary {
op: BinOp,
lhs: NodeId,
rhs: NodeId,
},
Unary {
op: UnOp,
expr: NodeId,
},
Array {
elements: Vec<NodeId>,
},
Table {
fields: IndexMap<String, NodeId>,
},
Quantity {
value: NodeId,
unit: Unit,
},
ApplyConstraint {
target: NodeId,
constraint: Constraint,
},
PatchDirective {
op: PatchOp,
},
SchemaField {
field: SchemaField,
},
SchemaTable {
fields: IndexMap<String, NodeId>,
closed: bool,
},
SchemaArray {
element: NodeId,
},
Semantic {
tag: SemanticTag,
inner: NodeId,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SemanticTag {
Prompt,
Policy,
Template,
}
impl SemanticTag {
pub fn as_str(&self) -> &'static str {
match self {
SemanticTag::Prompt => "prompt",
SemanticTag::Policy => "policy",
SemanticTag::Template => "template",
}
}
pub fn parse_tag(s: &str) -> Option<Self> {
match s {
"prompt" => Some(SemanticTag::Prompt),
"policy" => Some(SemanticTag::Policy),
"template" => Some(SemanticTag::Template),
_ => None,
}
}
}
#[derive(Debug, Clone)]
pub enum InterpPart {
Literal(String),
Expr(NodeId),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BinOp {
Add,
Sub,
Mul,
Div,
Eq,
Ne,
Lt,
Le,
Gt,
Ge,
And,
Or,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum UnOp {
Neg,
Not,
}
#[derive(Clone, Deserialize)]
pub enum Value {
Null,
Bool(bool),
Int(i64),
Float(f64),
String(String),
DurationMs(u64), Quantity(f64, Unit), Array(Vec<Value>),
Table(IndexMap<String, Value>),
Secret(String), }
impl std::fmt::Debug for Value {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Value::Secret(_) => write!(f, "Secret(***)"),
Value::Array(items) => {
write!(f, "[")?;
for (i, item) in items.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(f, "{item:?}")?;
}
write!(f, "]")
}
Value::Table(map) => {
write!(f, "{{")?;
for (i, (k, val)) in map.iter().enumerate() {
if i > 0 {
write!(f, ", ")?;
}
write!(f, "{k:?}: {val:?}")?;
}
write!(f, "}}")
}
Value::Null => write!(f, "Null"),
Value::Bool(b) => write!(f, "{b:?}"),
Value::Int(i) => write!(f, "{i:?}"),
Value::Float(x) => write!(f, "{x:?}"),
Value::String(s) => write!(f, "{s:?}"),
Value::DurationMs(ms) => write!(f, "DurationMs({ms})"),
Value::Quantity(n, unit) => write!(f, "Quantity({n:?}, {unit:?})"),
}
}
}
impl Serialize for Value {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
let json = match self {
Value::Null => serde_json::Value::Null,
Value::Bool(b) => serde_json::Value::from(*b),
Value::Int(i) => serde_json::Value::from(*i),
Value::Float(x) => serde_json::Value::from(*x),
Value::String(s) => serde_json::Value::from(s.clone()),
Value::DurationMs(ms) => serde_json::json!({ "ms": ms }),
Value::Quantity(n, unit) => serde_json::json!({ unit.suffix(): n }),
Value::Array(items) => serde_json::Value::Array(
items
.iter()
.map(|v| serde_json::to_value(v).unwrap_or(serde_json::Value::Null))
.collect(),
),
Value::Table(map) => {
let mut obj = serde_json::Map::new();
for (k, val) in map {
obj.insert(
k.clone(),
serde_json::to_value(val).unwrap_or(serde_json::Value::Null),
);
}
serde_json::Value::Object(obj)
}
Value::Secret(_) => serde_json::Value::from("***SECRET***"),
};
json.serialize(serializer)
}
}
impl Value {
pub fn as_str(&self) -> Option<&str> {
match self {
Value::String(s) => Some(s),
_ => None,
}
}
pub fn as_int(&self) -> Option<i64> {
match self {
Value::Int(i) => Some(*i),
_ => None,
}
}
}
#[derive(Debug, Clone)]
pub enum PatchOp {
Set { path: String, value: NodeId },
Remove { path: String },
Rename { from: String, to: String },
}