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    // --- Module and Symbols ---
14    Symbol(String),
15    /// Import a symbol from another module: (ModuleName, SymbolName)
16    Import(String, String),
17    /// Export a symbol: (SymbolName, Body)
18    Export(String, Id),
19    /// A module container: (ModuleName, Items)
20    Module(String, Vec<Id>),
21
22    // --- Basic Intents ---
23    Map(Id, Id),
24    Filter(Id, Id),
25    Reduce(Id, Id, Id),
26    StateUpdate(Id, Id),
27    Choice(Id, Id, Id),
28    Repeat(Id, Id),
29    /// LifeCycle: (setup, cleanup)
30    LifeCycle(Id, Id),
31    /// Meta: (body) for compile-time execution
32    Meta(Id),
33    /// Trap: (body) for error handling/interruption
34    Trap(Id),
35    /// Return: (value) for returning from a function
36    Return(Id),
37
38    // --- Composition ---
39    Seq(Vec<Id>),
40    Compose(Id, Id),
41
42    // --- Context and Constraints ---
43    WithContext(Id, Id),
44    WithConstraint(Id, Id),
45
46    // --- Context Atoms ---
47    CpuContext,
48    GpuContext,
49    AsyncContext,
50    SpatialContext,
51    ComptimeContext,
52    ResourceContext,
53    SafeContext,
54
55    // --- Constraint Atoms ---
56    EffectConstraint(crate::constraint::Effect),
57    OwnershipConstraint(crate::constraint::Ownership),
58    TypeConstraint(String),
59    AtomicConstraint,
60
61    // --- Extension Point ---
62    Extension(String, Vec<Id>),
63
64    // --- Cross Language ---
65    /// Cross-language call
66    CrossLangCall(CrossLanguageCall),
67
68    // --- Concretization Layer (Phase 5) ---
69    SoALayout(Id),
70    AoSLayout(Id),
71    Tiled(usize, Id),
72    Unrolled(usize, Id),
73    Vectorized(usize, Id),
74
75    // --- Optimized Specialized Intents (Avoid E-Graph cycles) ---
76    /// Tiled version of a Map: (size, f, x)
77    TiledMap(usize, Id, Id),
78    /// Vectorized version of a Map: (width, f, x)
79    VectorizedMap(usize, Id, Id),
80    /// Unrolled version of a Map: (factor, f, x)
81    UnrolledMap(usize, Id, Id),
82    /// SoA Layout applied to a Map: (f, x)
83    SoAMap(Id, Id),
84
85    // 特化的扩展，方便在规则中使用
86    GpuMap(Id, Id),
87    CpuMap(Id, Id),
88
89    // --- Spatial Layer (Phase 7) ---
90    /// Declare a pipeline stage: (body, metadata)
91    Pipe(Id, Id),
92    /// Explicit register: (value)
93    Reg(Id),
94
95    // --- Function and Closure (New) ---
96    /// Lambda definition: (params, body)
97    Lambda(Vec<String>, Id),
98    /// Function application: (func, args)
99    Apply(Id, Vec<Id>),
100    /// Closure capture: (body, captured_vars)
101    Closure(Id, Vec<Id>),
102
103    // --- Resource Management ---
104    /// Explicit Reference Count Clone: (target)
105    ResourceClone(Id),
106    /// Explicit Reference Count Drop: (target)
107    ResourceDrop(Id),
108}
109
110#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
111pub struct CrossLanguageCall {
112    /// The target language/platform
113    pub language: String,
114    /// Full name of module/libiary
115    pub module_path: String,
116    /// Unique function name
117    pub function_name: String,
118    /// The fixed arguments
119    pub arguments: Vec<Id>,
120}
121
122impl Language for IKun {
123    fn children(&self) -> Vec<Id> {
124        match self {
125            IKun::Constant(_)
126            | IKun::FloatConstant(_)
127            | IKun::BooleanConstant(_)
128            | IKun::StringConstant(_)
129            | IKun::Symbol(_)
130            | IKun::CpuContext
131            | IKun::GpuContext
132            | IKun::AsyncContext
133            | IKun::SpatialContext
134            | IKun::ComptimeContext
135            | IKun::ResourceContext
136            | IKun::SafeContext
137            | IKun::EffectConstraint(_)
138            | IKun::OwnershipConstraint(_)
139            | IKun::TypeConstraint(_)
140            | IKun::AtomicConstraint
141            | IKun::Import(_, _) => vec![],
142
143            IKun::Map(a, b)
144            | IKun::Filter(a, b)
145            | IKun::StateUpdate(a, b)
146            | IKun::Repeat(a, b)
147            | IKun::LifeCycle(a, b)
148            | IKun::Compose(a, b)
149            | IKun::WithContext(a, b)
150            | IKun::WithConstraint(a, b)
151            | IKun::GpuMap(a, b)
152            | IKun::CpuMap(a, b)
153            | IKun::Pipe(a, b)
154            | IKun::SoAMap(a, b)
155            | IKun::TiledMap(_, a, b)
156            | IKun::VectorizedMap(_, a, b)
157            | IKun::UnrolledMap(_, a, b) => vec![*a, *b],
158
159            IKun::Reduce(a, b, c) | IKun::Choice(a, b, c) => vec![*a, *b, *c],
160
161            IKun::Meta(a)
162            | IKun::Trap(a)
163            | IKun::Return(a)
164            | IKun::Reg(a)
165            | IKun::SoALayout(a)
166            | IKun::AoSLayout(a)
167            | IKun::Tiled(_, a)
168            | IKun::Unrolled(_, a)
169            | IKun::Vectorized(_, a)
170            | IKun::Lambda(_, a)
171            | IKun::ResourceClone(a)
172            | IKun::ResourceDrop(a)
173            | IKun::Export(_, a) => vec![*a],
174
175            IKun::Seq(ids) | IKun::Extension(_, ids) | IKun::Module(_, ids) => ids.clone(),
176
177            IKun::CrossLangCall(call) => call.arguments.clone(),
178
179            IKun::Apply(a, ids) | IKun::Closure(a, ids) => {
180                let mut res = vec![*a];
181                res.extend(ids.iter().cloned());
182                res
183            }
184        }
185    }
186
187    fn map_children(&self, mut f: impl FnMut(Id) -> Id) -> Self {
188        match self {
189            IKun::Constant(v) => IKun::Constant(*v),
190            IKun::FloatConstant(v) => IKun::FloatConstant(*v),
191            IKun::BooleanConstant(v) => IKun::BooleanConstant(*v),
192            IKun::StringConstant(s) => IKun::StringConstant(s.clone()),
193            IKun::Symbol(s) => IKun::Symbol(s.clone()),
194            IKun::Map(a, b) => IKun::Map(f(*a), f(*b)),
195            IKun::Filter(a, b) => IKun::Filter(f(*a), f(*b)),
196            IKun::Reduce(a, b, c) => IKun::Reduce(f(*a), f(*b), f(*c)),
197            IKun::StateUpdate(a, b) => IKun::StateUpdate(f(*a), f(*b)),
198            IKun::Choice(a, b, c) => IKun::Choice(f(*a), f(*b), f(*c)),
199            IKun::Repeat(a, b) => IKun::Repeat(f(*a), f(*b)),
200            IKun::LifeCycle(a, b) => IKun::LifeCycle(f(*a), f(*b)),
201            IKun::Meta(a) => IKun::Meta(f(*a)),
202            IKun::Trap(a) => IKun::Trap(f(*a)),
203            IKun::Return(a) => IKun::Return(f(*a)),
204            IKun::Seq(ids) => IKun::Seq(ids.iter().map(|id| f(*id)).collect()),
205            IKun::Compose(a, b) => IKun::Compose(f(*a), f(*b)),
206            IKun::WithContext(a, b) => IKun::WithContext(f(*a), f(*b)),
207            IKun::WithConstraint(a, b) => IKun::WithConstraint(f(*a), f(*b)),
208            IKun::CpuContext => IKun::CpuContext,
209            IKun::GpuContext => IKun::GpuContext,
210            IKun::AsyncContext => IKun::AsyncContext,
211            IKun::SpatialContext => IKun::SpatialContext,
212            IKun::ComptimeContext => IKun::ComptimeContext,
213            IKun::ResourceContext => IKun::ResourceContext,
214            IKun::SafeContext => IKun::SafeContext,
215            IKun::EffectConstraint(e) => IKun::EffectConstraint(e.clone()),
216            IKun::OwnershipConstraint(o) => IKun::OwnershipConstraint(o.clone()),
217            IKun::TypeConstraint(t) => IKun::TypeConstraint(t.clone()),
218            IKun::AtomicConstraint => IKun::AtomicConstraint,
219            IKun::Extension(name, ids) => {
220                IKun::Extension(name.clone(), ids.iter().map(|id| f(*id)).collect())
221            }
222            IKun::CrossLangCall(call) => IKun::CrossLangCall(CrossLanguageCall {
223                language: call.language.clone(),
224                module_path: call.module_path.clone(),
225                function_name: call.function_name.clone(),
226                arguments: call.arguments.iter().map(|id| f(*id)).collect(),
227            }),
228            IKun::SoALayout(a) => IKun::SoALayout(f(*a)),
229            IKun::AoSLayout(a) => IKun::AoSLayout(f(*a)),
230            IKun::Tiled(s, a) => IKun::Tiled(*s, f(*a)),
231            IKun::Unrolled(s, a) => IKun::Unrolled(*s, f(*a)),
232            IKun::Vectorized(s, a) => IKun::Vectorized(*s, f(*a)),
233            IKun::TiledMap(s, a, b) => IKun::TiledMap(*s, f(*a), f(*b)),
234            IKun::VectorizedMap(s, a, b) => IKun::VectorizedMap(*s, f(*a), f(*b)),
235            IKun::UnrolledMap(s, a, b) => IKun::UnrolledMap(*s, f(*a), f(*b)),
236            IKun::SoAMap(a, b) => IKun::SoAMap(f(*a), f(*b)),
237            IKun::GpuMap(a, b) => IKun::GpuMap(f(*a), f(*b)),
238            IKun::CpuMap(a, b) => IKun::CpuMap(f(*a), f(*b)),
239            IKun::Pipe(a, b) => IKun::Pipe(f(*a), f(*b)),
240            IKun::Reg(a) => IKun::Reg(f(*a)),
241            IKun::Lambda(params, body) => IKun::Lambda(params.clone(), f(*body)),
242            IKun::Apply(func, args) => {
243                IKun::Apply(f(*func), args.iter().map(|id| f(*id)).collect())
244            }
245            IKun::Closure(body, captured) => {
246                IKun::Closure(f(*body), captured.iter().map(|id| f(*id)).collect())
247            }
248            IKun::ResourceClone(a) => IKun::ResourceClone(f(*a)),
249            IKun::ResourceDrop(a) => IKun::ResourceDrop(f(*a)),
250            IKun::Import(m, s) => IKun::Import(m.clone(), s.clone()),
251            IKun::Export(s, b) => IKun::Export(s.clone(), f(*b)),
252            IKun::Module(m, ids) => IKun::Module(m.clone(), ids.iter().map(|id| f(*id)).collect()),
253        }
254    }
255}
256
257impl fmt::Display for IKun {
258    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
259        match self {
260            IKun::Constant(v) => write!(f, "{}", v),
261            IKun::FloatConstant(v) => write!(f, "{}", f64::from_bits(*v)),
262            IKun::BooleanConstant(v) => write!(f, "{}", v),
263            IKun::StringConstant(s) => write!(f, "\"{}\"", s),
264            IKun::Symbol(s) => write!(f, "{}", s),
265            IKun::Map(_, _) => write!(f, "map"),
266            IKun::Filter(_, _) => write!(f, "filter"),
267            IKun::Reduce(_, _, _) => write!(f, "reduce"),
268            IKun::StateUpdate(_, _) => write!(f, "state-update"),
269            IKun::Choice(_, _, _) => write!(f, "choice"),
270            IKun::Repeat(_, _) => write!(f, "repeat"),
271            IKun::Seq(_) => write!(f, "seq"),
272            IKun::Compose(_, _) => write!(f, "compose"),
273            IKun::WithContext(_, _) => write!(f, "with-context"),
274            IKun::WithConstraint(_, _) => write!(f, "with-constraint"),
275            IKun::CpuContext => write!(f, "cpu-context"),
276            IKun::GpuContext => write!(f, "gpu-context"),
277            IKun::AsyncContext => write!(f, "async-context"),
278            IKun::SpatialContext => write!(f, "spatial-context"),
279            IKun::ComptimeContext => write!(f, "comptime-context"),
280            IKun::ResourceContext => write!(f, "res-context"),
281            IKun::SafeContext => write!(f, "safe-context"),
282            IKun::LifeCycle(_, _) => write!(f, "lifecycle"),
283            IKun::Meta(_) => write!(f, "meta"),
284            IKun::Trap(_) => write!(f, "trap"),
285            IKun::Return(_) => write!(f, "return"),
286            IKun::EffectConstraint(e) => write!(f, "effect:{:?}", e),
287            IKun::OwnershipConstraint(o) => write!(f, "ownership:{:?}", o),
288            IKun::TypeConstraint(t) => write!(f, "type:{}", t),
289            IKun::AtomicConstraint => write!(f, "atomic-constraint"),
290            IKun::Extension(name, _) => write!(f, "ext:{}", name),
291            IKun::CrossLangCall(call) => {
292                write!(
293                    f,
294                    "cross-call:{}:{}:{}",
295                    call.language, call.module_path, call.function_name
296                )
297            }
298            IKun::GpuMap(_, _) => write!(f, "gpu-map"),
299            IKun::CpuMap(_, _) => write!(f, "cpu-map"),
300            IKun::SoALayout(_) => write!(f, "soa"),
301            IKun::AoSLayout(_) => write!(f, "aos"),
302            IKun::Tiled(s, _) => write!(f, "tiled:{}", s),
303            IKun::Unrolled(s, _) => write!(f, "unrolled:{}", s),
304            IKun::Vectorized(s, _) => write!(f, "vectorized:{}", s),
305            IKun::SoAMap(_, _) => write!(f, "soa-map"),
306            IKun::TiledMap(s, _, _) => write!(f, "tiled-map:{}", s),
307            IKun::VectorizedMap(s, _, _) => write!(f, "vectorized-map:{}", s),
308            IKun::UnrolledMap(s, _, _) => write!(f, "unrolled-map:{}", s),
309            IKun::Pipe(_, _) => write!(f, "pipe"),
310            IKun::Reg(_) => write!(f, "reg"),
311            IKun::Lambda(params, _) => write!(f, "lambda({})", params.join(", ")),
312            IKun::Apply(_, _) => write!(f, "apply"),
313            IKun::Closure(_, _) => write!(f, "closure"),
314            IKun::ResourceClone(_) => write!(f, "res-clone"),
315            IKun::ResourceDrop(_) => write!(f, "res-drop"),
316            IKun::Import(m, s) => write!(f, "import:{}:{}", m, s),
317            IKun::Export(s, _) => write!(f, "export:{}", s),
318            IKun::Module(m, _) => write!(f, "module:{}", m),
319        }
320    }
321}
322
323pub type Intent = IKun;
324pub type IntentOp = IKun;
325
326/// Recursive tree structure for Backends
327#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
328pub enum IKunTree {
329    Constant(i64),
330    FloatConstant(u64), // Use bits for Eq/Hash
331    BooleanConstant(bool),
332    StringConstant(String),
333    Symbol(String),
334    Map(Box<IKunTree>, Box<IKunTree>),
335    Filter(Box<IKunTree>, Box<IKunTree>),
336    Reduce(Box<IKunTree>, Box<IKunTree>, Box<IKunTree>),
337    StateUpdate(Box<IKunTree>, Box<IKunTree>),
338    Choice(Box<IKunTree>, Box<IKunTree>, Box<IKunTree>),
339    Repeat(Box<IKunTree>, Box<IKunTree>),
340    LifeCycle(Box<IKunTree>, Box<IKunTree>),
341    Meta(Box<IKunTree>),
342    Trap(Box<IKunTree>),
343    Return(Box<IKunTree>),
344    Seq(Vec<IKunTree>),
345    Compose(Box<IKunTree>, Box<IKunTree>),
346    WithContext(Box<IKunTree>, Box<IKunTree>),
347    WithConstraint(Box<IKunTree>, Box<IKunTree>),
348    CpuContext,
349    GpuContext,
350    AsyncContext,
351    SpatialContext,
352    ComptimeContext,
353    ResourceContext,
354    SafeContext,
355    EffectConstraint(crate::constraint::Effect),
356    OwnershipConstraint(crate::constraint::Ownership),
357    TypeConstraint(String),
358    AtomicConstraint,
359    Extension(String, Vec<IKunTree>),
360    CrossLangCall {
361        language: String,
362        module_path: String,
363        function_name: String,
364        arguments: Vec<IKunTree>,
365    },
366    GpuMap(Box<IKunTree>, Box<IKunTree>),
367    CpuMap(Box<IKunTree>, Box<IKunTree>),
368    TiledMap(usize, Box<IKunTree>, Box<IKunTree>),
369    VectorizedMap(usize, Box<IKunTree>, Box<IKunTree>),
370    UnrolledMap(usize, Box<IKunTree>, Box<IKunTree>),
371    SoAMap(Box<IKunTree>, Box<IKunTree>),
372
373    // Phase 5: Concretization
374    SoALayout(Box<IKunTree>),
375    AoSLayout(Box<IKunTree>),
376    Tiled(usize, Box<IKunTree>),
377    Unrolled(usize, Box<IKunTree>),
378    Vectorized(usize, Box<IKunTree>),
379
380    // Phase 7: Spatial
381    Pipe(Box<IKunTree>, Box<IKunTree>),
382    Reg(Box<IKunTree>),
383
384    // --- Function and Closure (New) ---
385    Lambda(Vec<String>, Box<IKunTree>),
386    Apply(Box<IKunTree>, Vec<IKunTree>),
387    Closure(Box<IKunTree>, Vec<IKunTree>),
388
389    ResourceClone(Box<IKunTree>),
390    ResourceDrop(Box<IKunTree>),
391
392    Import(String, String),
393    Export(String, Box<IKunTree>),
394    Module(String, Vec<IKunTree>),
395
396    /// Source location information for debugging
397    Source(chomsky_types::Loc, Box<IKunTree>),
398}
399
400impl Default for IKunTree {
401    fn default() -> Self {
402        IKunTree::Module("default".to_string(), Vec::new())
403    }
404}
405
406impl IKunTree {
407    pub fn to_egraph<A: crate::egraph::Analysis<IKun>>(
408        &self,
409        egraph: &mut crate::egraph::EGraph<IKun, A>,
410    ) -> Id {
411        match self {
412            IKunTree::Constant(v) => egraph.add(IKun::Constant(*v)),
413            IKunTree::FloatConstant(v) => egraph.add(IKun::FloatConstant(*v)),
414            IKunTree::BooleanConstant(v) => egraph.add(IKun::BooleanConstant(*v)),
415            IKunTree::StringConstant(s) => egraph.add(IKun::StringConstant(s.clone())),
416            IKunTree::Symbol(s) => egraph.add(IKun::Symbol(s.clone())),
417            IKunTree::Map(f, x) => {
418                let f_id = f.to_egraph(egraph);
419                let x_id = x.to_egraph(egraph);
420                egraph.add(IKun::Map(f_id, x_id))
421            }
422            IKunTree::Filter(f, x) => {
423                let f_id = f.to_egraph(egraph);
424                let x_id = x.to_egraph(egraph);
425                egraph.add(IKun::Filter(f_id, x_id))
426            }
427            IKunTree::Reduce(f, init, list) => {
428                let f_id = f.to_egraph(egraph);
429                let init_id = init.to_egraph(egraph);
430                let list_id = list.to_egraph(egraph);
431                egraph.add(IKun::Reduce(f_id, init_id, list_id))
432            }
433            IKunTree::StateUpdate(var, val) => {
434                let var_id = var.to_egraph(egraph);
435                let val_id = val.to_egraph(egraph);
436                egraph.add(IKun::StateUpdate(var_id, val_id))
437            }
438            IKunTree::Choice(cond, t, f) => {
439                let cond_id = cond.to_egraph(egraph);
440                let t_id = t.to_egraph(egraph);
441                let f_id = f.to_egraph(egraph);
442                egraph.add(IKun::Choice(cond_id, t_id, f_id))
443            }
444            IKunTree::Repeat(count, body) => {
445                let count_id = count.to_egraph(egraph);
446                let body_id = body.to_egraph(egraph);
447                egraph.add(IKun::Repeat(count_id, body_id))
448            }
449            IKunTree::LifeCycle(setup, cleanup) => {
450                let setup_id = setup.to_egraph(egraph);
451                let cleanup_id = cleanup.to_egraph(egraph);
452                egraph.add(IKun::LifeCycle(setup_id, cleanup_id))
453            }
454            IKunTree::Meta(body) => {
455                let body_id = body.to_egraph(egraph);
456                egraph.add(IKun::Meta(body_id))
457            }
458            IKunTree::Trap(body) => {
459                let body_id = body.to_egraph(egraph);
460                egraph.add(IKun::Trap(body_id))
461            }
462            IKunTree::Return(val) => {
463                let val_id = val.to_egraph(egraph);
464                egraph.add(IKun::Return(val_id))
465            }
466            IKunTree::Seq(trees) => {
467                let ids = trees.iter().map(|t| t.to_egraph(egraph)).collect();
468                egraph.add(IKun::Seq(ids))
469            }
470            IKunTree::Compose(f, g) => {
471                let f_id = f.to_egraph(egraph);
472                let g_id = g.to_egraph(egraph);
473                egraph.add(IKun::Compose(f_id, g_id))
474            }
475            IKunTree::WithContext(ctx, body) => {
476                let ctx_id = ctx.to_egraph(egraph);
477                let body_id = body.to_egraph(egraph);
478                egraph.add(IKun::WithContext(ctx_id, body_id))
479            }
480            IKunTree::WithConstraint(c, body) => {
481                let c_id = c.to_egraph(egraph);
482                let body_id = body.to_egraph(egraph);
483                egraph.add(IKun::WithConstraint(c_id, body_id))
484            }
485            IKunTree::CpuContext => egraph.add(IKun::CpuContext),
486            IKunTree::GpuContext => egraph.add(IKun::GpuContext),
487            IKunTree::AsyncContext => egraph.add(IKun::AsyncContext),
488            IKunTree::SpatialContext => egraph.add(IKun::SpatialContext),
489            IKunTree::ComptimeContext => egraph.add(IKun::ComptimeContext),
490            IKunTree::ResourceContext => egraph.add(IKun::ResourceContext),
491            IKunTree::SafeContext => egraph.add(IKun::SafeContext),
492            IKunTree::EffectConstraint(e) => egraph.add(IKun::EffectConstraint(e.clone())),
493            IKunTree::OwnershipConstraint(o) => egraph.add(IKun::OwnershipConstraint(o.clone())),
494            IKunTree::TypeConstraint(t) => egraph.add(IKun::TypeConstraint(t.clone())),
495            IKunTree::AtomicConstraint => egraph.add(IKun::AtomicConstraint),
496            IKunTree::Extension(name, trees) => {
497                let ids = trees.iter().map(|t| t.to_egraph(egraph)).collect();
498                egraph.add(IKun::Extension(name.clone(), ids))
499            }
500            IKunTree::Source(_, body) => body.to_egraph(egraph),
501            IKunTree::GpuMap(f, x) => {
502                let f_id = f.to_egraph(egraph);
503                let x_id = x.to_egraph(egraph);
504                egraph.add(IKun::GpuMap(f_id, x_id))
505            }
506            IKunTree::CpuMap(f, x) => {
507                let f_id = f.to_egraph(egraph);
508                let x_id = x.to_egraph(egraph);
509                egraph.add(IKun::CpuMap(f_id, x_id))
510            }
511            IKunTree::TiledMap(s, f, x) => {
512                let f_id = f.to_egraph(egraph);
513                let x_id = x.to_egraph(egraph);
514                egraph.add(IKun::TiledMap(*s, f_id, x_id))
515            }
516            IKunTree::VectorizedMap(s, f, x) => {
517                let f_id = f.to_egraph(egraph);
518                let x_id = x.to_egraph(egraph);
519                egraph.add(IKun::VectorizedMap(*s, f_id, x_id))
520            }
521            IKunTree::UnrolledMap(s, f, x) => {
522                let f_id = f.to_egraph(egraph);
523                let x_id = x.to_egraph(egraph);
524                egraph.add(IKun::UnrolledMap(*s, f_id, x_id))
525            }
526            IKunTree::SoAMap(f, x) => {
527                let f_id = f.to_egraph(egraph);
528                let x_id = x.to_egraph(egraph);
529                egraph.add(IKun::SoAMap(f_id, x_id))
530            }
531            IKunTree::SoALayout(a) => {
532                let a_id = a.to_egraph(egraph);
533                egraph.add(IKun::SoALayout(a_id))
534            }
535            IKunTree::AoSLayout(a) => {
536                let a_id = a.to_egraph(egraph);
537                egraph.add(IKun::AoSLayout(a_id))
538            }
539            IKunTree::Tiled(s, a) => {
540                let a_id = a.to_egraph(egraph);
541                egraph.add(IKun::Tiled(*s, a_id))
542            }
543            IKunTree::Unrolled(s, a) => {
544                let a_id = a.to_egraph(egraph);
545                egraph.add(IKun::Unrolled(*s, a_id))
546            }
547            IKunTree::Vectorized(s, a) => {
548                let a_id = a.to_egraph(egraph);
549                egraph.add(IKun::Vectorized(*s, a_id))
550            }
551            IKunTree::Pipe(a, b) => {
552                let a_id = a.to_egraph(egraph);
553                let b_id = b.to_egraph(egraph);
554                egraph.add(IKun::Pipe(a_id, b_id))
555            }
556            IKunTree::Reg(a) => {
557                let a_id = a.to_egraph(egraph);
558                egraph.add(IKun::Reg(a_id))
559            }
560            IKunTree::Lambda(params, body) => {
561                let body_id = body.to_egraph(egraph);
562                egraph.add(IKun::Lambda(params.clone(), body_id))
563            }
564            IKunTree::Apply(func, args) => {
565                let func_id = func.to_egraph(egraph);
566                let arg_ids = args.iter().map(|a| a.to_egraph(egraph)).collect();
567                egraph.add(IKun::Apply(func_id, arg_ids))
568            }
569            IKunTree::Closure(body, captured) => {
570                let body_id = body.to_egraph(egraph);
571                let captured_ids = captured.iter().map(|c| c.to_egraph(egraph)).collect();
572                egraph.add(IKun::Closure(body_id, captured_ids))
573            }
574            IKunTree::ResourceClone(target) => {
575                let target_id = target.to_egraph(egraph);
576                egraph.add(IKun::ResourceClone(target_id))
577            }
578            IKunTree::ResourceDrop(target) => {
579                let target_id = target.to_egraph(egraph);
580                egraph.add(IKun::ResourceDrop(target_id))
581            }
582            IKunTree::CrossLangCall {
583                language,
584                module_path,
585                function_name,
586                arguments,
587            } => {
588                let args_ids = arguments.iter().map(|arg| arg.to_egraph(egraph)).collect();
589                egraph.add(IKun::CrossLangCall(CrossLanguageCall {
590                    language: language.clone(),
591                    module_path: module_path.clone(),
592                    function_name: function_name.clone(),
593                    arguments: args_ids,
594                }))
595            }
596            IKunTree::Import(m, s) => egraph.add(IKun::Import(m.clone(), s.clone())),
597            IKunTree::Export(s, b) => {
598                let b_id = b.to_egraph(egraph);
599                egraph.add(IKun::Export(s.clone(), b_id))
600            }
601            IKunTree::Module(m, trees) => {
602                let ids = trees.iter().map(|t| t.to_egraph(egraph)).collect();
603                egraph.add(IKun::Module(m.clone(), ids))
604            }
605        }
606    }
607}
chomsky_uir/intent.rs

chomsky_uir/
intent.rs
