chomsky_uir/
intent.rs

1use crate::egraph::Language;
2use crate::union_find::Id;
3use serde::{Deserialize, Serialize};
4use std::fmt;
5
6#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
7pub enum IKun {
8    // --- Data Atoms ---
9    Constant(i64),
10    FloatConstant(u64), // Use bits for Eq/Hash/Ord
11    BooleanConstant(bool),
12    StringConstant(String),
13    /// None/Nil/Null constant
14    None,
15    // --- Module and Symbols ---
16    Symbol(String),
17    /// Import a symbol from another module: (ModuleName, SymbolName)
18    Import(String, String),
19    /// Export a symbol: (SymbolName, Body)
20    Export(String, Id),
21    /// A module container: (ModuleName, Items)
22    Module(String, Vec<Id>),
23
24    // --- Basic Intents ---
25    Map(Id, Id),
26    Filter(Id, Id),
27    Reduce(Id, Id, Id),
28    /// State update: (key, value)
29    StateUpdate(Id, Id),
30    Choice(Id, Id, Id),
31    Repeat(Id, Id),
32    /// LifeCycle: (setup, cleanup)
33    LifeCycle(Id, Id),
34    /// Meta: (body) for compile-time execution
35    Meta(Id),
36    /// Trap: (body) for error handling/interruption
37    Trap(Id),
38    /// Return: (value) for returning from a function
39    Return(Id),
40
41    // --- Composition ---
42    Seq(Vec<Id>),
43    Compose(Id, Id),
44
45    // --- Context and Constraints ---
46    WithContext(Id, Id),
47    WithConstraint(Id, Id),
48
49    // --- Context Atoms ---
50    CpuContext,
51    GpuContext,
52    AsyncContext,
53    SpatialContext,
54    ComptimeContext,
55    ResourceContext,
56    SafeContext,
57
58    // --- Constraint Atoms ---
59    EffectConstraint(crate::constraint::Effect),
60    OwnershipConstraint(crate::constraint::Ownership),
61    TypeConstraint(String),
62    AtomicConstraint,
63
64    // --- Extension Point ---
65    Extension(String, Vec<Id>),
66
67    // --- Cross Language ---
68    /// Cross-language call
69    CrossLangCall(CrossLanguageCall),
70
71    // --- Concretization Layer (Phase 5) ---
72    SoALayout(Id),
73    AoSLayout(Id),
74    Tiled(usize, Id),
75    Unrolled(usize, Id),
76    Vectorized(usize, Id),
77
78    // --- Optimized Specialized Intents (Avoid E-Graph cycles) ---
79    /// Tiled version of a Map: (size, f, x)
80    TiledMap(usize, Id, Id),
81    /// Vectorized version of a Map: (width, f, x)
82    VectorizedMap(usize, Id, Id),
83    /// Unrolled version of a Map: (factor, f, x)
84    UnrolledMap(usize, Id, Id),
85    /// SoA Layout applied to a Map: (f, x)
86    SoAMap(Id, Id),
87
88    // 特化的扩展，方便在规则中使用
89    GpuMap(Id, Id),
90    CpuMap(Id, Id),
91
92    // --- Spatial Layer (Phase 7) ---
93    /// Declare a pipeline stage: (body, metadata)
94    Pipe(Id, Id),
95    /// Explicit register: (value)
96    Reg(Id),
97
98    // --- Function and Closure (New) ---
99    /// Lambda definition: (params, body)
100    Lambda(Vec<String>, Id),
101    /// Function application: (func, args)
102    Apply(Id, Vec<Id>),
103    /// Closure capture: (body, captured_vars)
104    Closure(Id, Vec<Id>),
105
106    // --- Resource Management ---
107    /// Explicit Reference Count Clone: (target)
108    ResourceClone(Id),
109    /// Explicit Reference Count Drop: (target)
110    ResourceDrop(Id),
111
112    // --- Pointer Operations ---
113    /// Take address of a symbol/location: (target)
114    AddrOf(Id),
115    /// Dereference a pointer: (pointer)
116    Deref(Id),
117    /// Pointer arithmetic: (base, offset)
118    PtrOffset(Id, Id),
119
120    // --- Object Model Access ---
121    /// Class definition: (name, bases, body)
122    ClassDef(String, Vec<Id>, Id),
123    /// Table/Object initialization: (pairs)
124    Table(Vec<Id>),
125    /// Key-value pair for table: (key, value)
126    Pair(Id, Id),
127    /// Get index: (receiver, index)
128    GetIndex(Id, Id),
129    /// Set index: (receiver, index, value)
130    SetIndex(Id, Id, Id),
131
132    // --- Operations ---
133    /// Binary operation: (op_name, left, right)
134    BinaryOp(String, Id, Id),
135    /// Unary operation: (op_name, operand)
136    UnaryOp(String, Id),
137
138    // --- Object Model Access (Legacy/Specialized) ---
139    /// Static access: (receiver, offset)
140    StaticAccess(Id, usize),
141    /// Witness access: (receiver, witness_table, index)
142    WitnessAccess(Id, Id, usize),
143    /// Dynamic access: (receiver, symbol_name, call_site_id)
144    DynamicAccess(Id, String, u32),
145
146    // --- Object Model Call ---
147    /// Static call: (function, args)
148    StaticCall(Id, Vec<Id>),
149    /// Witness call: (receiver, witness_table, index, args)
150    WitnessCall(Id, Id, usize, Vec<Id>),
151    /// Dynamic call: (receiver, symbol_name, args, call_site_id)
152    DynamicCall(Id, String, Vec<Id>, u32),
153}
154
155#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
156pub struct CrossLanguageCall {
157    /// The target language/platform
158    pub language: String,
159    /// Full name of module/libiary
160    pub module_path: String,
161    /// Unique function name
162    pub function_name: String,
163    /// The fixed arguments
164    pub arguments: Vec<Id>,
165}
166
167impl Language for IKun {
168    fn children(&self) -> Vec<Id> {
169        match self {
170            IKun::Constant(_)
171            | IKun::FloatConstant(_)
172            | IKun::BooleanConstant(_)
173            | IKun::StringConstant(_)
174            | IKun::None
175            | IKun::Symbol(_)
176            | IKun::CpuContext
177            | IKun::GpuContext
178            | IKun::AsyncContext
179            | IKun::SpatialContext
180            | IKun::ComptimeContext
181            | IKun::ResourceContext
182            | IKun::SafeContext
183            | IKun::EffectConstraint(_)
184            | IKun::OwnershipConstraint(_)
185            | IKun::TypeConstraint(_)
186            | IKun::AtomicConstraint
187            | IKun::Import(_, _) => vec![],
188
189            IKun::Map(a, b)
190            | IKun::Filter(a, b)
191            | IKun::StateUpdate(a, b)
192            | IKun::Repeat(a, b)
193            | IKun::LifeCycle(a, b)
194            | IKun::Compose(a, b)
195            | IKun::WithContext(a, b)
196            | IKun::WithConstraint(a, b)
197            | IKun::GpuMap(a, b)
198            | IKun::CpuMap(a, b)
199            | IKun::Pipe(a, b)
200            | IKun::SoAMap(a, b)
201            | IKun::TiledMap(_, a, b)
202            | IKun::VectorizedMap(_, a, b)
203            | IKun::UnrolledMap(_, a, b) => vec![*a, *b],
204
205            IKun::Reduce(a, b, c) | IKun::Choice(a, b, c) => vec![*a, *b, *c],
206
207            IKun::Meta(a)
208            | IKun::Trap(a)
209            | IKun::Return(a)
210            | IKun::Reg(a)
211            | IKun::SoALayout(a)
212            | IKun::AoSLayout(a)
213            | IKun::Tiled(_, a)
214            | IKun::Unrolled(_, a)
215            | IKun::Vectorized(_, a)
216            | IKun::Lambda(_, a)
217            | IKun::ResourceClone(a)
218            | IKun::ResourceDrop(a)
219            | IKun::AddrOf(a)
220            | IKun::Deref(a)
221            | IKun::UnaryOp(_, a)
222            | IKun::StaticAccess(a, _)
223            | IKun::DynamicAccess(a, _, _)
224            | IKun::Export(_, a) => vec![*a],
225
226            IKun::WitnessAccess(a, b, _)
227            | IKun::PtrOffset(a, b)
228            | IKun::Pair(a, b)
229            | IKun::GetIndex(a, b)
230            | IKun::BinaryOp(_, a, b) => vec![*a, *b],
231
232            IKun::SetIndex(a, b, c) => vec![*a, *b, *c],
233
234            IKun::Seq(ids) | IKun::Extension(_, ids) | IKun::Module(_, ids) | IKun::Table(ids) => {
235                ids.clone()
236            }
237
238            IKun::CrossLangCall(call) => call.arguments.clone(),
239
240            IKun::ClassDef(_, bases, body) => {
241                let mut res = bases.clone();
242                res.push(*body);
243                res
244            }
245
246            IKun::StaticCall(a, ids) | IKun::Apply(a, ids) | IKun::Closure(a, ids) => {
247                let mut res = vec![*a];
248                res.extend(ids.iter().cloned());
249                res
250            }
251            IKun::WitnessCall(a, b, _, ids) => {
252                let mut res = vec![*a, *b];
253                res.extend(ids.iter().cloned());
254                res
255            }
256            IKun::DynamicCall(a, _, ids, _) => {
257                let mut res = vec![*a];
258                res.extend(ids.iter().cloned());
259                res
260            }
261        }
262    }
263
264    fn map_children(&self, mut f: impl FnMut(Id) -> Id) -> Self {
265        match self {
266            IKun::Constant(v) => IKun::Constant(*v),
267            IKun::FloatConstant(v) => IKun::FloatConstant(*v),
268            IKun::BooleanConstant(v) => IKun::BooleanConstant(*v),
269            IKun::StringConstant(s) => IKun::StringConstant(s.clone()),
270            IKun::None => IKun::None,
271            IKun::Symbol(s) => IKun::Symbol(s.clone()),
272            IKun::Map(a, b) => IKun::Map(f(*a), f(*b)),
273            IKun::Filter(a, b) => IKun::Filter(f(*a), f(*b)),
274            IKun::Reduce(a, b, c) => IKun::Reduce(f(*a), f(*b), f(*c)),
275            IKun::StateUpdate(a, b) => IKun::StateUpdate(f(*a), f(*b)),
276            IKun::Choice(a, b, c) => IKun::Choice(f(*a), f(*b), f(*c)),
277            IKun::Repeat(a, b) => IKun::Repeat(f(*a), f(*b)),
278            IKun::LifeCycle(a, b) => IKun::LifeCycle(f(*a), f(*b)),
279            IKun::Meta(a) => IKun::Meta(f(*a)),
280            IKun::Trap(a) => IKun::Trap(f(*a)),
281            IKun::Return(a) => IKun::Return(f(*a)),
282            IKun::Seq(ids) => IKun::Seq(ids.iter().map(|id| f(*id)).collect()),
283            IKun::Compose(a, b) => IKun::Compose(f(*a), f(*b)),
284            IKun::WithContext(a, b) => IKun::WithContext(f(*a), f(*b)),
285            IKun::WithConstraint(a, b) => IKun::WithConstraint(f(*a), f(*b)),
286            IKun::CpuContext => IKun::CpuContext,
287            IKun::GpuContext => IKun::GpuContext,
288            IKun::AsyncContext => IKun::AsyncContext,
289            IKun::SpatialContext => IKun::SpatialContext,
290            IKun::ComptimeContext => IKun::ComptimeContext,
291            IKun::ResourceContext => IKun::ResourceContext,
292            IKun::SafeContext => IKun::SafeContext,
293            IKun::EffectConstraint(e) => IKun::EffectConstraint(e.clone()),
294            IKun::OwnershipConstraint(o) => IKun::OwnershipConstraint(o.clone()),
295            IKun::TypeConstraint(t) => IKun::TypeConstraint(t.clone()),
296            IKun::AtomicConstraint => IKun::AtomicConstraint,
297            IKun::Extension(name, ids) => {
298                IKun::Extension(name.clone(), ids.iter().map(|id| f(*id)).collect())
299            }
300            IKun::CrossLangCall(call) => IKun::CrossLangCall(CrossLanguageCall {
301                language: call.language.clone(),
302                module_path: call.module_path.clone(),
303                function_name: call.function_name.clone(),
304                arguments: call.arguments.iter().map(|id| f(*id)).collect(),
305            }),
306            IKun::SoALayout(a) => IKun::SoALayout(f(*a)),
307            IKun::AoSLayout(a) => IKun::AoSLayout(f(*a)),
308            IKun::Tiled(s, a) => IKun::Tiled(*s, f(*a)),
309            IKun::Unrolled(s, a) => IKun::Unrolled(*s, f(*a)),
310            IKun::Vectorized(s, a) => IKun::Vectorized(*s, f(*a)),
311            IKun::TiledMap(s, a, b) => IKun::TiledMap(*s, f(*a), f(*b)),
312            IKun::VectorizedMap(s, a, b) => IKun::VectorizedMap(*s, f(*a), f(*b)),
313            IKun::UnrolledMap(s, a, b) => IKun::UnrolledMap(*s, f(*a), f(*b)),
314            IKun::SoAMap(a, b) => IKun::SoAMap(f(*a), f(*b)),
315            IKun::GpuMap(a, b) => IKun::GpuMap(f(*a), f(*b)),
316            IKun::CpuMap(a, b) => IKun::CpuMap(f(*a), f(*b)),
317            IKun::Pipe(a, b) => IKun::Pipe(f(*a), f(*b)),
318            IKun::Reg(a) => IKun::Reg(f(*a)),
319            IKun::Lambda(params, body) => IKun::Lambda(params.clone(), f(*body)),
320            IKun::Apply(func, args) => {
321                IKun::Apply(f(*func), args.iter().map(|id| f(*id)).collect())
322            }
323            IKun::Closure(body, captured) => {
324                IKun::Closure(f(*body), captured.iter().map(|id| f(*id)).collect())
325            }
326            IKun::ResourceClone(a) => IKun::ResourceClone(f(*a)),
327            IKun::ResourceDrop(a) => IKun::ResourceDrop(f(*a)),
328            IKun::AddrOf(a) => IKun::AddrOf(f(*a)),
329            IKun::Deref(a) => IKun::Deref(f(*a)),
330            IKun::PtrOffset(a, b) => IKun::PtrOffset(f(*a), f(*b)),
331            IKun::ClassDef(name, bases, body) => IKun::ClassDef(
332                name.clone(),
333                bases.iter().map(|id| f(*id)).collect(),
334                f(*body),
335            ),
336            IKun::Table(ids) => IKun::Table(ids.iter().map(|id| f(*id)).collect()),
337            IKun::Pair(a, b) => IKun::Pair(f(*a), f(*b)),
338            IKun::GetIndex(a, b) => IKun::GetIndex(f(*a), f(*b)),
339            IKun::SetIndex(a, b, c) => IKun::SetIndex(f(*a), f(*b), f(*c)),
340            IKun::BinaryOp(op, a, b) => IKun::BinaryOp(op.clone(), f(*a), f(*b)),
341            IKun::UnaryOp(op, a) => IKun::UnaryOp(op.clone(), f(*a)),
342            IKun::Import(m, s) => IKun::Import(m.clone(), s.clone()),
343            IKun::Export(s, b) => IKun::Export(s.clone(), f(*b)),
344            IKun::Module(m, ids) => IKun::Module(m.clone(), ids.iter().map(|id| f(*id)).collect()),
345            IKun::StaticAccess(a, s) => IKun::StaticAccess(f(*a), *s),
346            IKun::WitnessAccess(a, b, s) => IKun::WitnessAccess(f(*a), f(*b), *s),
347            IKun::DynamicAccess(a, s, id) => IKun::DynamicAccess(f(*a), s.clone(), *id),
348            IKun::StaticCall(a, ids) => {
349                IKun::StaticCall(f(*a), ids.iter().map(|id| f(*id)).collect())
350            }
351            IKun::WitnessCall(a, b, s, ids) => {
352                IKun::WitnessCall(f(*a), f(*b), *s, ids.iter().map(|id| f(*id)).collect())
353            }
354            IKun::DynamicCall(a, s, ids, id) => {
355                IKun::DynamicCall(f(*a), s.clone(), ids.iter().map(|id| f(*id)).collect(), *id)
356            }
357        }
358    }
359}
360
361impl fmt::Display for IKun {
362    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
363        match self {
364            IKun::Constant(v) => write!(f, "{}", v),
365            IKun::FloatConstant(v) => write!(f, "{}", f64::from_bits(*v)),
366            IKun::BooleanConstant(v) => write!(f, "{}", v),
367            IKun::StringConstant(s) => write!(f, "\"{}\"", s),
368            IKun::None => write!(f, "none"),
369            IKun::Symbol(s) => write!(f, "{}", s),
370            IKun::Map(_, _) => write!(f, "map"),
371            IKun::Filter(_, _) => write!(f, "filter"),
372            IKun::Reduce(_, _, _) => write!(f, "reduce"),
373            IKun::StateUpdate(_, _) => write!(f, "state-update"),
374            IKun::Choice(_, _, _) => write!(f, "choice"),
375            IKun::Repeat(_, _) => write!(f, "repeat"),
376            IKun::Seq(_) => write!(f, "seq"),
377            IKun::Compose(_, _) => write!(f, "compose"),
378            IKun::WithContext(_, _) => write!(f, "with-context"),
379            IKun::WithConstraint(_, _) => write!(f, "with-constraint"),
380            IKun::CpuContext => write!(f, "cpu-context"),
381            IKun::GpuContext => write!(f, "gpu-context"),
382            IKun::AsyncContext => write!(f, "async-context"),
383            IKun::SpatialContext => write!(f, "spatial-context"),
384            IKun::ComptimeContext => write!(f, "comptime-context"),
385            IKun::ResourceContext => write!(f, "res-context"),
386            IKun::SafeContext => write!(f, "safe-context"),
387            IKun::LifeCycle(_, _) => write!(f, "lifecycle"),
388            IKun::Meta(_) => write!(f, "meta"),
389            IKun::Trap(_) => write!(f, "trap"),
390            IKun::Return(_) => write!(f, "return"),
391            IKun::EffectConstraint(e) => write!(f, "effect:{:?}", e),
392            IKun::OwnershipConstraint(o) => write!(f, "ownership:{:?}", o),
393            IKun::TypeConstraint(t) => write!(f, "type:{}", t),
394            IKun::AtomicConstraint => write!(f, "atomic-constraint"),
395            IKun::Extension(name, _) => write!(f, "ext:{}", name),
396            IKun::CrossLangCall(call) => {
397                write!(
398                    f,
399                    "cross-call:{}:{}:{}",
400                    call.language, call.module_path, call.function_name
401                )
402            }
403            IKun::GpuMap(_, _) => write!(f, "gpu-map"),
404            IKun::CpuMap(_, _) => write!(f, "cpu-map"),
405            IKun::SoALayout(_) => write!(f, "soa"),
406            IKun::AoSLayout(_) => write!(f, "aos"),
407            IKun::Tiled(s, _) => write!(f, "tiled:{}", s),
408            IKun::Unrolled(s, _) => write!(f, "unrolled:{}", s),
409            IKun::Vectorized(s, _) => write!(f, "vectorized:{}", s),
410            IKun::SoAMap(_, _) => write!(f, "soa-map"),
411            IKun::TiledMap(s, _, _) => write!(f, "tiled-map:{}", s),
412            IKun::VectorizedMap(s, _, _) => write!(f, "vectorized-map:{}", s),
413            IKun::UnrolledMap(s, _, _) => write!(f, "unrolled-map:{}", s),
414            IKun::Pipe(_, _) => write!(f, "pipe"),
415            IKun::Reg(_) => write!(f, "reg"),
416            IKun::Lambda(params, _) => write!(f, "lambda({})", params.join(", ")),
417            IKun::Apply(_, _) => write!(f, "apply"),
418            IKun::Closure(_, _) => write!(f, "closure"),
419            IKun::ResourceClone(_) => write!(f, "res-clone"),
420            IKun::ResourceDrop(_) => write!(f, "res-drop"),
421            IKun::AddrOf(_) => write!(f, "address-of"),
422            IKun::Deref(_) => write!(f, "deref"),
423            IKun::PtrOffset(_, _) => write!(f, "ptr-offset"),
424            IKun::ClassDef(name, _, _) => write!(f, "class:{}", name),
425            IKun::Table(_) => write!(f, "table"),
426            IKun::Pair(_, _) => write!(f, "pair"),
427            IKun::GetIndex(_, _) => write!(f, "get-index"),
428            IKun::SetIndex(_, _, _) => write!(f, "set-index"),
429            IKun::BinaryOp(op, _, _) => write!(f, "bin-op:{}", op),
430            IKun::UnaryOp(op, _) => write!(f, "unary-op:{}", op),
431            IKun::Import(m, s) => write!(f, "import:{}:{}", m, s),
432            IKun::Export(s, _) => write!(f, "export:{}", s),
433            IKun::Module(m, _) => write!(f, "module:{}", m),
434            IKun::StaticAccess(_, s) => write!(f, "static-access:{}", s),
435            IKun::WitnessAccess(_, _, s) => write!(f, "witness-access:{}", s),
436            IKun::DynamicAccess(_, s, id) => write!(f, "dyn-access:{}@{}", s, id),
437            IKun::StaticCall(_, _) => write!(f, "static-call"),
438            IKun::WitnessCall(_, _, s, _) => write!(f, "witness-call:{}", s),
439            IKun::DynamicCall(_, s, _, _) => write!(f, "dyn-call:{}", s),
440        }
441    }
442}
443
444pub type Intent = IKun;
445pub type IntentOp = IKun;
446
447/// Recursive tree structure for Backends
448#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
449pub enum IKunTree {
450    Constant(i64),
451    FloatConstant(u64), // Use bits for Eq/Hash
452    BooleanConstant(bool),
453    StringConstant(String),
454    None,
455    Symbol(String),
456    Map(Box<IKunTree>, Box<IKunTree>),
457    Filter(Box<IKunTree>, Box<IKunTree>),
458    Reduce(Box<IKunTree>, Box<IKunTree>, Box<IKunTree>),
459    StateUpdate(Box<IKunTree>, Box<IKunTree>),
460    Choice(Box<IKunTree>, Box<IKunTree>, Box<IKunTree>),
461    Repeat(Box<IKunTree>, Box<IKunTree>),
462    LifeCycle(Box<IKunTree>, Box<IKunTree>),
463    Meta(Box<IKunTree>),
464    Trap(Box<IKunTree>),
465    Return(Box<IKunTree>),
466    Seq(Vec<IKunTree>),
467    Compose(Box<IKunTree>, Box<IKunTree>),
468    WithContext(Box<IKunTree>, Box<IKunTree>),
469    WithConstraint(Box<IKunTree>, Box<IKunTree>),
470    CpuContext,
471    GpuContext,
472    AsyncContext,
473    SpatialContext,
474    ComptimeContext,
475    ResourceContext,
476    SafeContext,
477    EffectConstraint(crate::constraint::Effect),
478    OwnershipConstraint(crate::constraint::Ownership),
479    TypeConstraint(String),
480    AtomicConstraint,
481    Extension(String, Vec<IKunTree>),
482    CrossLangCall {
483        language: String,
484        module_path: String,
485        function_name: String,
486        arguments: Vec<IKunTree>,
487    },
488    GpuMap(Box<IKunTree>, Box<IKunTree>),
489    CpuMap(Box<IKunTree>, Box<IKunTree>),
490    TiledMap(usize, Box<IKunTree>, Box<IKunTree>),
491    VectorizedMap(usize, Box<IKunTree>, Box<IKunTree>),
492    UnrolledMap(usize, Box<IKunTree>, Box<IKunTree>),
493    SoAMap(Box<IKunTree>, Box<IKunTree>),
494
495    // Phase 5: Concretization
496    SoALayout(Box<IKunTree>),
497    AoSLayout(Box<IKunTree>),
498    Tiled(usize, Box<IKunTree>),
499    Unrolled(usize, Box<IKunTree>),
500    Vectorized(usize, Box<IKunTree>),
501
502    // Phase 7: Spatial
503    Pipe(Box<IKunTree>, Box<IKunTree>),
504    Reg(Box<IKunTree>),
505
506    // --- Function and Closure (New) ---
507    Lambda(Vec<String>, Box<IKunTree>),
508    Apply(Box<IKunTree>, Vec<IKunTree>),
509    Closure(Box<IKunTree>, Vec<IKunTree>),
510
511    ResourceClone(Box<IKunTree>),
512    ResourceDrop(Box<IKunTree>),
513
514    // --- Pointer Operations ---
515    AddrOf(Box<IKunTree>),
516    Deref(Box<IKunTree>),
517    PtrOffset(Box<IKunTree>, Box<IKunTree>),
518
519    // --- Object Model ---
520    ClassDef(String, Vec<IKunTree>, Box<IKunTree>),
521    Table(Vec<IKunTree>),
522    Pair(Box<IKunTree>, Box<IKunTree>),
523    GetIndex(Box<IKunTree>, Box<IKunTree>),
524    SetIndex(Box<IKunTree>, Box<IKunTree>, Box<IKunTree>),
525
526    // --- Operations ---
527    BinaryOp(String, Box<IKunTree>, Box<IKunTree>),
528    UnaryOp(String, Box<IKunTree>),
529
530    Import(String, String),
531    Export(String, Box<IKunTree>),
532    Module(String, Vec<IKunTree>),
533
534    // --- Object Model Access ---
535    /// Static access: (receiver, offset)
536    StaticAccess(Box<IKunTree>, usize),
537    /// Witness access: (receiver, witness_table, index)
538    WitnessAccess(Box<IKunTree>, Box<IKunTree>, usize),
539    /// Dynamic access: (receiver, symbol_name, call_site_id)
540    DynamicAccess(Box<IKunTree>, String, u32),
541
542    // --- Object Model Call ---
543    /// Static call: (function, args)
544    StaticCall(Box<IKunTree>, Vec<IKunTree>),
545    /// Witness call: (receiver, witness_table, index, args)
546    WitnessCall(Box<IKunTree>, Box<IKunTree>, usize, Vec<IKunTree>),
547    /// Dynamic call: (receiver, symbol_name, args, call_site_id)
548    DynamicCall(Box<IKunTree>, String, Vec<IKunTree>, u32),
549
550    /// Source location information for debugging
551    Source(chomsky_types::Loc, Box<IKunTree>),
552}
553
554impl Default for IKunTree {
555    fn default() -> Self {
556        IKunTree::Module("default".to_string(), Vec::new())
557    }
558}
559
560impl IKunTree {
561    pub fn to_egraph<A: crate::egraph::Analysis<IKun>>(
562        &self,
563        egraph: &mut crate::egraph::EGraph<IKun, A>,
564    ) -> Id {
565        match self {
566            IKunTree::Constant(v) => egraph.add(IKun::Constant(*v)),
567            IKunTree::FloatConstant(v) => egraph.add(IKun::FloatConstant(*v)),
568            IKunTree::BooleanConstant(v) => egraph.add(IKun::BooleanConstant(*v)),
569            IKunTree::StringConstant(s) => egraph.add(IKun::StringConstant(s.clone())),
570            IKunTree::None => egraph.add(IKun::None),
571            IKunTree::Symbol(s) => egraph.add(IKun::Symbol(s.clone())),
572            IKunTree::Map(f, x) => {
573                let f_id = f.to_egraph(egraph);
574                let x_id = x.to_egraph(egraph);
575                egraph.add(IKun::Map(f_id, x_id))
576            }
577            IKunTree::Filter(f, x) => {
578                let f_id = f.to_egraph(egraph);
579                let x_id = x.to_egraph(egraph);
580                egraph.add(IKun::Filter(f_id, x_id))
581            }
582            IKunTree::Reduce(f, init, list) => {
583                let f_id = f.to_egraph(egraph);
584                let init_id = init.to_egraph(egraph);
585                let list_id = list.to_egraph(egraph);
586                egraph.add(IKun::Reduce(f_id, init_id, list_id))
587            }
588            IKunTree::StateUpdate(var, val) => {
589                let var_id = var.to_egraph(egraph);
590                let val_id = val.to_egraph(egraph);
591                egraph.add(IKun::StateUpdate(var_id, val_id))
592            }
593            IKunTree::Choice(cond, t, f) => {
594                let cond_id = cond.to_egraph(egraph);
595                let t_id = t.to_egraph(egraph);
596                let f_id = f.to_egraph(egraph);
597                egraph.add(IKun::Choice(cond_id, t_id, f_id))
598            }
599            IKunTree::Repeat(count, body) => {
600                let count_id = count.to_egraph(egraph);
601                let body_id = body.to_egraph(egraph);
602                egraph.add(IKun::Repeat(count_id, body_id))
603            }
604            IKunTree::LifeCycle(setup, cleanup) => {
605                let setup_id = setup.to_egraph(egraph);
606                let cleanup_id = cleanup.to_egraph(egraph);
607                egraph.add(IKun::LifeCycle(setup_id, cleanup_id))
608            }
609            IKunTree::Meta(body) => {
610                let body_id = body.to_egraph(egraph);
611                egraph.add(IKun::Meta(body_id))
612            }
613            IKunTree::Trap(body) => {
614                let body_id = body.to_egraph(egraph);
615                egraph.add(IKun::Trap(body_id))
616            }
617            IKunTree::Return(val) => {
618                let val_id = val.to_egraph(egraph);
619                egraph.add(IKun::Return(val_id))
620            }
621            IKunTree::Seq(trees) => {
622                let ids = trees.iter().map(|t| t.to_egraph(egraph)).collect();
623                egraph.add(IKun::Seq(ids))
624            }
625            IKunTree::Compose(f, g) => {
626                let f_id = f.to_egraph(egraph);
627                let g_id = g.to_egraph(egraph);
628                egraph.add(IKun::Compose(f_id, g_id))
629            }
630            IKunTree::WithContext(ctx, body) => {
631                let ctx_id = ctx.to_egraph(egraph);
632                let body_id = body.to_egraph(egraph);
633                egraph.add(IKun::WithContext(ctx_id, body_id))
634            }
635            IKunTree::WithConstraint(c, body) => {
636                let c_id = c.to_egraph(egraph);
637                let body_id = body.to_egraph(egraph);
638                egraph.add(IKun::WithConstraint(c_id, body_id))
639            }
640            IKunTree::CpuContext => egraph.add(IKun::CpuContext),
641            IKunTree::GpuContext => egraph.add(IKun::GpuContext),
642            IKunTree::AsyncContext => egraph.add(IKun::AsyncContext),
643            IKunTree::SpatialContext => egraph.add(IKun::SpatialContext),
644            IKunTree::ComptimeContext => egraph.add(IKun::ComptimeContext),
645            IKunTree::ResourceContext => egraph.add(IKun::ResourceContext),
646            IKunTree::SafeContext => egraph.add(IKun::SafeContext),
647            IKunTree::EffectConstraint(e) => egraph.add(IKun::EffectConstraint(e.clone())),
648            IKunTree::OwnershipConstraint(o) => egraph.add(IKun::OwnershipConstraint(o.clone())),
649            IKunTree::TypeConstraint(t) => egraph.add(IKun::TypeConstraint(t.clone())),
650            IKunTree::AtomicConstraint => egraph.add(IKun::AtomicConstraint),
651            IKunTree::Extension(name, trees) => {
652                let ids = trees.iter().map(|t| t.to_egraph(egraph)).collect();
653                egraph.add(IKun::Extension(name.clone(), ids))
654            }
655            IKunTree::Source(_, body) => body.to_egraph(egraph),
656            IKunTree::GpuMap(f, x) => {
657                let f_id = f.to_egraph(egraph);
658                let x_id = x.to_egraph(egraph);
659                egraph.add(IKun::GpuMap(f_id, x_id))
660            }
661            IKunTree::CpuMap(f, x) => {
662                let f_id = f.to_egraph(egraph);
663                let x_id = x.to_egraph(egraph);
664                egraph.add(IKun::CpuMap(f_id, x_id))
665            }
666            IKunTree::TiledMap(s, f, x) => {
667                let f_id = f.to_egraph(egraph);
668                let x_id = x.to_egraph(egraph);
669                egraph.add(IKun::TiledMap(*s, f_id, x_id))
670            }
671            IKunTree::VectorizedMap(s, f, x) => {
672                let f_id = f.to_egraph(egraph);
673                let x_id = x.to_egraph(egraph);
674                egraph.add(IKun::VectorizedMap(*s, f_id, x_id))
675            }
676            IKunTree::UnrolledMap(s, f, x) => {
677                let f_id = f.to_egraph(egraph);
678                let x_id = x.to_egraph(egraph);
679                egraph.add(IKun::UnrolledMap(*s, f_id, x_id))
680            }
681            IKunTree::SoAMap(f, x) => {
682                let f_id = f.to_egraph(egraph);
683                let x_id = x.to_egraph(egraph);
684                egraph.add(IKun::SoAMap(f_id, x_id))
685            }
686            IKunTree::SoALayout(a) => {
687                let a_id = a.to_egraph(egraph);
688                egraph.add(IKun::SoALayout(a_id))
689            }
690            IKunTree::AoSLayout(a) => {
691                let a_id = a.to_egraph(egraph);
692                egraph.add(IKun::AoSLayout(a_id))
693            }
694            IKunTree::Tiled(s, a) => {
695                let a_id = a.to_egraph(egraph);
696                egraph.add(IKun::Tiled(*s, a_id))
697            }
698            IKunTree::Unrolled(s, a) => {
699                let a_id = a.to_egraph(egraph);
700                egraph.add(IKun::Unrolled(*s, a_id))
701            }
702            IKunTree::Vectorized(s, a) => {
703                let a_id = a.to_egraph(egraph);
704                egraph.add(IKun::Vectorized(*s, a_id))
705            }
706            IKunTree::Pipe(a, b) => {
707                let a_id = a.to_egraph(egraph);
708                let b_id = b.to_egraph(egraph);
709                egraph.add(IKun::Pipe(a_id, b_id))
710            }
711            IKunTree::Reg(a) => {
712                let a_id = a.to_egraph(egraph);
713                egraph.add(IKun::Reg(a_id))
714            }
715            IKunTree::Lambda(params, body) => {
716                let body_id = body.to_egraph(egraph);
717                egraph.add(IKun::Lambda(params.clone(), body_id))
718            }
719            IKunTree::Apply(func, args) => {
720                let func_id = func.to_egraph(egraph);
721                let arg_ids = args.iter().map(|a| a.to_egraph(egraph)).collect();
722                egraph.add(IKun::Apply(func_id, arg_ids))
723            }
724            IKunTree::Closure(body, captured) => {
725                let body_id = body.to_egraph(egraph);
726                let captured_ids = captured.iter().map(|c| c.to_egraph(egraph)).collect();
727                egraph.add(IKun::Closure(body_id, captured_ids))
728            }
729            IKunTree::ResourceClone(target) => {
730                let target_id = target.to_egraph(egraph);
731                egraph.add(IKun::ResourceClone(target_id))
732            }
733            IKunTree::ResourceDrop(target) => {
734                let target_id = target.to_egraph(egraph);
735                egraph.add(IKun::ResourceDrop(target_id))
736            }
737            IKunTree::CrossLangCall {
738                language,
739                module_path,
740                function_name,
741                arguments,
742            } => {
743                let args_ids = arguments.iter().map(|arg| arg.to_egraph(egraph)).collect();
744                egraph.add(IKun::CrossLangCall(CrossLanguageCall {
745                    language: language.clone(),
746                    module_path: module_path.clone(),
747                    function_name: function_name.clone(),
748                    arguments: args_ids,
749                }))
750            }
751            IKunTree::AddrOf(a) => {
752                let a_id = a.to_egraph(egraph);
753                egraph.add(IKun::AddrOf(a_id))
754            }
755            IKunTree::Deref(a) => {
756                let a_id = a.to_egraph(egraph);
757                egraph.add(IKun::Deref(a_id))
758            }
759            IKunTree::PtrOffset(a, b) => {
760                let a_id = a.to_egraph(egraph);
761                let b_id = b.to_egraph(egraph);
762                egraph.add(IKun::PtrOffset(a_id, b_id))
763            }
764            IKunTree::ClassDef(name, bases, body) => {
765                let bases_ids = bases.iter().map(|b| b.to_egraph(egraph)).collect();
766                let body_id = body.to_egraph(egraph);
767                egraph.add(IKun::ClassDef(name.clone(), bases_ids, body_id))
768            }
769            IKunTree::Table(ids) => {
770                let table_ids = ids.iter().map(|id| id.to_egraph(egraph)).collect();
771                egraph.add(IKun::Table(table_ids))
772            }
773            IKunTree::Pair(a, b) => {
774                let a_id = a.to_egraph(egraph);
775                let b_id = b.to_egraph(egraph);
776                egraph.add(IKun::Pair(a_id, b_id))
777            }
778            IKunTree::GetIndex(a, b) => {
779                let a_id = a.to_egraph(egraph);
780                let b_id = b.to_egraph(egraph);
781                egraph.add(IKun::GetIndex(a_id, b_id))
782            }
783            IKunTree::SetIndex(a, b, c) => {
784                let a_id = a.to_egraph(egraph);
785                let b_id = b.to_egraph(egraph);
786                let c_id = c.to_egraph(egraph);
787                egraph.add(IKun::SetIndex(a_id, b_id, c_id))
788            }
789            IKunTree::BinaryOp(op, a, b) => {
790                let a_id = a.to_egraph(egraph);
791                let b_id = b.to_egraph(egraph);
792                egraph.add(IKun::BinaryOp(op.clone(), a_id, b_id))
793            }
794            IKunTree::UnaryOp(op, a) => {
795                let a_id = a.to_egraph(egraph);
796                egraph.add(IKun::UnaryOp(op.clone(), a_id))
797            }
798            IKunTree::Import(m, s) => egraph.add(IKun::Import(m.clone(), s.clone())),
799            IKunTree::Export(s, b) => {
800                let b_id = b.to_egraph(egraph);
801                egraph.add(IKun::Export(s.clone(), b_id))
802            }
803            IKunTree::Module(m, trees) => {
804                let ids = trees.iter().map(|t| t.to_egraph(egraph)).collect();
805                egraph.add(IKun::Module(m.clone(), ids))
806            }
807            IKunTree::StaticAccess(a, s) => {
808                let a_id = a.to_egraph(egraph);
809                egraph.add(IKun::StaticAccess(a_id, *s))
810            }
811            IKunTree::WitnessAccess(a, b, s) => {
812                let a_id = a.to_egraph(egraph);
813                let b_id = b.to_egraph(egraph);
814                egraph.add(IKun::WitnessAccess(a_id, b_id, *s))
815            }
816            IKunTree::DynamicAccess(a, s, id) => {
817                let a_id = a.to_egraph(egraph);
818                egraph.add(IKun::DynamicAccess(a_id, s.clone(), *id))
819            }
820            IKunTree::StaticCall(a, ids) => {
821                let a_id = a.to_egraph(egraph);
822                let ids = ids.iter().map(|id| id.to_egraph(egraph)).collect();
823                egraph.add(IKun::StaticCall(a_id, ids))
824            }
825            IKunTree::WitnessCall(receiver, witness_table, index, args) => {
826                let receiver_id = receiver.to_egraph(egraph);
827                let witness_table_id = witness_table.to_egraph(egraph);
828                let arg_ids = args.iter().map(|arg| arg.to_egraph(egraph)).collect();
829                egraph.add(IKun::WitnessCall(
830                    receiver_id,
831                    witness_table_id,
832                    *index,
833                    arg_ids,
834                ))
835            }
836            IKunTree::DynamicCall(receiver, method, args, call_site_id) => {
837                let receiver_id = receiver.to_egraph(egraph);
838                let arg_ids = args.iter().map(|arg| arg.to_egraph(egraph)).collect();
839                egraph.add(IKun::DynamicCall(
840                    receiver_id,
841                    method.clone(),
842                    arg_ids,
843                    *call_site_id,
844                ))
845            }
846        }
847    }
848}
chomsky_uir/intent.rs

chomsky_uir/
intent.rs
