use std::{
collections::hash_map::DefaultHasher,
hash::{Hash, Hasher},
sync::atomic::AtomicU64,
sync::Arc,
sync::OnceLock,
};
use indexmap::IndexMap;
use crate::parse::ast::*;
pub use crate::builtins::BuiltinMethod;
#[derive(Debug, Clone)]
pub struct CompiledCall {
pub method: BuiltinMethod,
pub name: Arc<str>,
pub sub_progs: Arc<[Arc<Program>]>,
pub sub_kernels: Arc<[crate::exec::pipeline::BodyKernel]>,
pub orig_args: Arc<[Arg]>,
pub demand_max_keep: Option<usize>,
}
#[derive(Debug, Clone)]
pub enum CompiledObjEntry {
Short {
name: Arc<str>,
ic: Arc<AtomicU64>,
},
Kv {
key: Arc<str>,
prog: Arc<Program>,
optional: bool,
cond: Option<Arc<Program>>,
},
KvPath {
key: Arc<str>,
steps: Arc<[KvStep]>,
optional: bool,
ics: Arc<[AtomicU64]>,
},
Dynamic {
key: Arc<Program>,
val: Arc<Program>,
},
Spread(Arc<Program>),
SpreadDeep(Arc<Program>),
}
#[derive(Debug, Clone)]
pub enum KvStep {
Field(Arc<str>),
Index(i64),
}
#[derive(Debug, Clone)]
pub enum CompiledFSPart {
Lit(Arc<str>),
Interp {
prog: Arc<Program>,
fmt: Option<FmtSpec>,
},
}
#[derive(Debug, Clone)]
pub struct BindObjSpec {
pub fields: Arc<[Arc<str>]>,
pub rest: Option<Arc<str>>,
}
#[derive(Debug, Clone)]
pub enum CompiledPipeStep {
Forward(Arc<Program>),
BindName(Arc<str>),
BindObj(Arc<BindObjSpec>),
BindArr(Arc<[Arc<str>]>),
}
#[derive(Debug, Clone)]
pub struct CompSpec {
pub expr: Arc<Program>,
pub vars: Arc<[Arc<str>]>,
pub iter: Arc<Program>,
pub cond: Option<Arc<Program>>,
}
#[derive(Debug, Clone)]
pub struct DictCompSpec {
pub key: Arc<Program>,
pub val: Arc<Program>,
pub vars: Arc<[Arc<str>]>,
pub iter: Arc<Program>,
pub cond: Option<Arc<Program>>,
}
#[derive(Debug, Clone)]
pub enum Opcode {
PushNull,
PushBool(bool),
PushInt(i64),
PushFloat(f64),
PushStr(Arc<str>),
PushRoot,
PushCurrent,
GetField(Arc<str>),
GetIndex(i64),
GetSlice(Option<i64>, Option<i64>),
DynIndex(Arc<Program>),
OptField(Arc<str>),
Descendant(Arc<str>),
DescendAll,
InlineFilter(Arc<Program>),
Quantifier(QuantifierKind),
RootChain(Arc<[Arc<str>]>),
FieldChain(Arc<FieldChainData>),
LoadIdent(Arc<str>),
Add,
Sub,
Mul,
Div,
Mod,
Eq,
Neq,
Lt,
Lte,
Gt,
Gte,
Fuzzy,
Not,
Neg,
CastOp(crate::parse::ast::CastType),
AndOp(Arc<Program>),
OrOp(Arc<Program>),
CoalesceOp(Arc<Program>),
CallMethod(Arc<CompiledCall>),
CallOptMethod(Arc<CompiledCall>),
MakeObj(Arc<[CompiledObjEntry]>),
MakeArr(Arc<[(Arc<Program>, bool)]>),
FString(Arc<[CompiledFSPart]>),
KindCheck {
ty: KindType,
negate: bool,
},
SetCurrent,
BindLamCurrent {
name: Option<Arc<str>>,
body: Arc<Program>,
},
PipelineRun {
base: Arc<Program>,
steps: Arc<[CompiledPipeStep]>,
},
LetExpr {
name: Arc<str>,
body: Arc<Program>,
},
IfElse {
then_: Arc<Program>,
else_: Arc<Program>,
},
TryExpr {
body: Arc<Program>,
default: Arc<Program>,
},
ListComp(Arc<CompSpec>),
DictComp(Arc<DictCompSpec>),
SetComp(Arc<CompSpec>),
PatchEval(Arc<CompiledPatch>),
DeleteMarkErr,
Match(Arc<CompiledMatch>),
DeepMatchAll(Arc<CompiledMatch>),
DeepMatchFirst(Arc<CompiledMatch>),
}
#[derive(Debug, Clone)]
pub enum MatchScrutinee {
Current,
Root,
Program(Arc<Program>),
}
#[derive(Debug, Clone)]
pub struct CompiledMatch {
pub scrutinee: MatchScrutinee,
pub ops: Arc<[MatchOp]>,
pub lits: Arc<[crate::parse::ast::PatLit]>,
pub guards: Arc<[Arc<Program>]>,
pub bodies: Arc<[Arc<Program>]>,
pub subpats: Arc<[crate::parse::ast::Pat]>,
pub max_slots: u16,
pub is_exhaustive: bool,
pub shape_summary: Option<MatchShapeSummary>,
}
#[derive(Debug, Clone)]
pub enum MatchShapeSummary {
ObjAnyOfKeys(Arc<[Arc<str>]>),
KindOnly(crate::parse::ast::KindType),
NumericRange {
lo: f64,
hi: f64,
inclusive: bool,
},
}
pub type MatchSlot = u16;
#[derive(Debug, Clone)]
pub enum MatchOp {
ResetArm {
slots: u16,
keep_above: u16,
},
KindCheck {
slot: MatchSlot,
kind: crate::parse::ast::KindType,
else_pc: u32,
},
LitEq {
slot: MatchSlot,
lit: u16,
else_pc: u32,
},
RangeCheck {
slot: MatchSlot,
lo: f64,
hi: f64,
inclusive: bool,
else_pc: u32,
},
ObjCheck {
slot: MatchSlot,
else_pc: u32,
},
LoadField {
src: MatchSlot,
key: Arc<str>,
dst: MatchSlot,
else_pc: u32,
},
LenCheck {
slot: MatchSlot,
len: u32,
exact: bool,
else_pc: u32,
},
LoadIndex {
src: MatchSlot,
idx: u32,
dst: MatchSlot,
},
LoadTail {
src: MatchSlot,
from: u32,
dst: MatchSlot,
},
LoadObjRest {
src: MatchSlot,
listed_keys: Arc<[Arc<str>]>,
dst: MatchSlot,
},
TestSubPat {
slot: MatchSlot,
subpat: u16,
else_pc: u32,
},
Bind {
name: Arc<str>,
slot: MatchSlot,
},
Guard {
prog: u16,
else_pc: u32,
},
Body {
prog: u16,
},
Fail,
Jump {
target_pc: u32,
},
}
#[derive(Debug, Clone)]
pub struct Program {
pub ops: Arc<[Opcode]>,
pub source: Arc<str>,
pub id: u64,
pub is_structural: bool,
pub ics: Arc<[AtomicU64]>,
}
#[derive(Debug)]
pub struct CompiledPatch {
pub root_prog: Arc<Program>,
pub ops: Vec<CompiledPatchOp>,
pub trie: OnceLock<Option<CompiledPatchTrie>>,
}
impl Clone for CompiledPatch {
fn clone(&self) -> Self {
Self {
root_prog: Arc::clone(&self.root_prog),
ops: self.ops.clone(),
trie: OnceLock::new(),
}
}
}
#[derive(Debug, Clone)]
pub struct CompiledPatchOp {
pub path: Vec<CompiledPathStep>,
pub val: CompiledPatchVal,
pub cond: Option<Arc<Program>>,
}
#[derive(Debug, Clone)]
pub enum CompiledPatchVal {
Replace(Arc<Program>),
Delete,
}
#[derive(Debug, Clone)]
pub enum CompiledPathStep {
Field(Arc<str>),
Index(i64),
DynIndex(Arc<Program>),
Wildcard,
WildcardFilter(Arc<Program>),
Descendant(Arc<str>),
}
#[derive(Debug, Clone)]
pub struct CompiledPatchTrie {
pub root: TrieNode,
}
#[derive(Debug, Clone)]
pub enum TrieNode {
Replace(Arc<Program>),
Delete,
Conditional {
cond_prog: Arc<Program>,
then: Box<TrieNode>,
},
Branch {
fields: IndexMap<Arc<str>, TrieNode>,
indices: Vec<(IdxKey, TrieNode)>,
},
}
#[derive(Debug, Clone)]
pub enum IdxKey {
Static(i64),
Dynamic(Arc<Program>),
}
impl CompiledPatchTrie {
pub fn from_ops(ops: &[CompiledPatchOp]) -> Option<Self> {
for op in ops {
for step in &op.path {
match step {
CompiledPathStep::Field(_)
| CompiledPathStep::Index(_)
| CompiledPathStep::DynIndex(_) => {}
CompiledPathStep::Wildcard
| CompiledPathStep::WildcardFilter(_)
| CompiledPathStep::Descendant(_) => return None,
}
}
}
let mut root = TrieNode::Branch {
fields: IndexMap::new(),
indices: Vec::new(),
};
for op in ops {
let leaf = match &op.val {
CompiledPatchVal::Replace(prog) => TrieNode::Replace(Arc::clone(prog)),
CompiledPatchVal::Delete => TrieNode::Delete,
};
let leaf = if let Some(cond_prog) = &op.cond {
TrieNode::Conditional {
cond_prog: Arc::clone(cond_prog),
then: Box::new(leaf),
}
} else {
leaf
};
insert_leaf(&mut root, &op.path, leaf);
}
Some(Self { root })
}
}
fn insert_leaf(node: &mut TrieNode, path: &[CompiledPathStep], leaf: TrieNode) {
if path.is_empty() {
*node = leaf;
return;
}
if !matches!(node, TrieNode::Branch { .. }) {
*node = TrieNode::Branch {
fields: IndexMap::new(),
indices: Vec::new(),
};
}
let TrieNode::Branch { fields, indices } = node else {
unreachable!()
};
match &path[0] {
CompiledPathStep::Field(name) => {
if let Some(child) = fields.get_mut(name) {
insert_leaf(child, &path[1..], leaf);
} else {
let mut fresh = TrieNode::Branch {
fields: IndexMap::new(),
indices: Vec::new(),
};
insert_leaf(&mut fresh, &path[1..], leaf);
fields.insert(Arc::clone(name), fresh);
}
}
CompiledPathStep::Index(i) => {
if let Some((_, child)) = indices
.iter_mut()
.find(|(k, _)| matches!(k, IdxKey::Static(s) if *s == *i))
{
insert_leaf(child, &path[1..], leaf);
} else {
let mut fresh = TrieNode::Branch {
fields: IndexMap::new(),
indices: Vec::new(),
};
insert_leaf(&mut fresh, &path[1..], leaf);
indices.push((IdxKey::Static(*i), fresh));
}
}
CompiledPathStep::DynIndex(prog) => {
if let Some((_, child)) = indices.iter_mut().find(|(k, _)| match k {
IdxKey::Dynamic(p) => Arc::ptr_eq(p, prog),
_ => false,
}) {
insert_leaf(child, &path[1..], leaf);
} else {
let mut fresh = TrieNode::Branch {
fields: IndexMap::new(),
indices: Vec::new(),
};
insert_leaf(&mut fresh, &path[1..], leaf);
indices.push((IdxKey::Dynamic(Arc::clone(prog)), fresh));
}
}
CompiledPathStep::Wildcard
| CompiledPathStep::WildcardFilter(_)
| CompiledPathStep::Descendant(_) => unreachable!(),
}
}
impl Program {
pub fn new(ops: Vec<Opcode>, source: &str) -> Self {
let id = hash_str(source);
let is_structural = ops.iter().all(|op| {
matches!(
op,
Opcode::PushRoot
| Opcode::PushCurrent
| Opcode::GetField(_)
| Opcode::GetIndex(_)
| Opcode::GetSlice(..)
| Opcode::OptField(_)
| Opcode::RootChain(_)
| Opcode::FieldChain(_)
)
});
let ics = fresh_ics(ops.len());
Self {
ops: ops.into(),
source: source.into(),
id,
is_structural,
ics,
}
}
}
#[derive(Debug)]
pub struct FieldChainData {
pub keys: Arc<[Arc<str>]>,
pub ics: Box<[AtomicU64]>,
}
impl FieldChainData {
pub fn new(keys: Arc<[Arc<str>]>) -> Self {
let n = keys.len();
let mut ics = Vec::with_capacity(n);
for _ in 0..n {
ics.push(AtomicU64::new(0));
}
Self {
keys,
ics: ics.into_boxed_slice(),
}
}
}
impl std::ops::Deref for FieldChainData {
type Target = [Arc<str>];
#[inline]
fn deref(&self) -> &[Arc<str>] {
&self.keys
}
}
pub fn fresh_ics(len: usize) -> Arc<[AtomicU64]> {
let mut v = Vec::with_capacity(len);
for _ in 0..len {
v.push(AtomicU64::new(0));
}
v.into()
}
#[inline]
pub(crate) fn disable_opcode_fusion() -> bool {
std::env::var_os("JETRO_DISABLE_OPCODE_FUSION").is_some()
}
fn hash_str(s: &str) -> u64 {
let mut h = DefaultHasher::new();
s.hash(&mut h);
h.finish()
}