portal_ir/
lib.rs

1use std::{
2    cell::RefCell,
3    collections::{BTreeMap, BTreeSet},
4    rc::Rc,
5};
6
7use id_arena::{Arena, Id};
8
9pub struct Value<T, Y, R, D> {
10    pub id: Id<ValueDef<T, Y, R, D>>,
11    pub idx: Option<u32>,
12}
13
14impl<T, Y, R, D> Clone for Value<T, Y, R, D> {
15    fn clone(&self) -> Self {
16        Self {
17            id: self.id.clone(),
18            idx: self.idx.clone(),
19        }
20    }
21}
22
23pub enum Param<T, Y, R, D> {
24    Value(Value<T, Y, R, D>),
25    Data(Id<D>),
26    Fun(Id<Fun<T, Y, R, D>>),
27}
28impl<T, Y, R, D> Clone for Param<T, Y, R, D> {
29    fn clone(&self) -> Self {
30        match self {
31            Self::Value(arg0) => Self::Value(arg0.clone()),
32            Self::Data(arg0) => Self::Data(arg0.clone()),
33            Self::Fun(arg0) => Self::Fun(arg0.clone()),
34        }
35    }
36}
37pub struct Params<T, Y, R, D> {
38    pub values: Vec<Value<T, Y, R, D>>,
39    pub data: Vec<Id<D>>,
40    pub funs: Vec<Id<Fun<T, Y, R, D>>>,
41}
42impl<T, Y, R, D> Default for Params<T, Y, R, D> {
43    fn default() -> Self {
44        Self {
45            values: Default::default(),
46            data: Default::default(),
47            funs: Default::default(),
48        }
49    }
50}
51impl<T, Y, R, D> Clone for Params<T, Y, R, D> {
52    fn clone(&self) -> Self {
53        Self {
54            values: self.values.clone(),
55            data: self.data.clone(),
56            funs: self.funs.clone(),
57        }
58    }
59}
60impl<T, Y, R, D> FromIterator<Param<T, Y, R, D>> for Params<T, Y, R, D> {
61    fn from_iter<I: IntoIterator<Item = Param<T, Y, R, D>>>(iter: I) -> Self {
62        // let mut i = iter.into_iter();
63        let mut p: Params<T, Y, R, D> = Default::default();
64        for j in iter.into_iter() {
65            match j {
66                Param::Value(v) => p.values.push(v),
67                Param::Data(d) => p.data.push(d),
68                Param::Fun(f) => p.funs.push(f),
69            }
70        }
71        return p;
72    }
73}
74impl<T: 'static, Y: 'static, R: 'static, D: 'static> IntoIterator for Params<T, Y, R, D> {
75    type Item = Param<T, Y, R, D>;
76
77    type IntoIter = Box<dyn Iterator<Item = Param<T, Y, R, D>>>;
78
79    fn into_iter(self) -> Self::IntoIter {
80        return Box::new(
81            self.values
82                .into_iter()
83                .map(Param::Value)
84                .chain(self.data.into_iter().map(Param::Data))
85                .chain(self.funs.into_iter().map(Param::Fun)),
86        );
87    }
88}
89
90pub enum ValueDef<T, Y, R, D> {
91    Param(usize),
92    Emit {
93        op: T,
94        params: Params<T, Y, R, D>,
95        typ: Y,
96        after: BTreeSet<Id<ValueDef<T, Y, R, D>>>,
97    },
98    Alias(Value<T, Y, R, D>),
99}
100
101impl<T: Clone, Y: Clone, R: Clone, D> Clone for ValueDef<T, Y, R, D> {
102    fn clone(&self) -> Self {
103        match self {
104            Self::Param(arg0) => Self::Param(arg0.clone()),
105            // Self::Emit(arg0, arg1, arg2,arg3) => Self::Emit(arg0.clone(), arg1.clone(), arg2.clone(),arg3.clone()),
106            Self::Alias(arg0) => Self::Alias(arg0.clone()),
107            // ValueDef::Param(_) => todo!(),
108            ValueDef::Emit {
109                op,
110                params,
111                typ,
112                after,
113            } => ValueDef::Emit {
114                op: op.clone(),
115                params: params.clone(),
116                typ: typ.clone(),
117                after: after.clone(),
118            },
119            // ValueDef::Alias(_) => todo!(),
120        }
121    }
122}
123
124pub struct Fun<T, Y, R, D> {
125    pub values: Arena<ValueDef<T, Y, R, D>>,
126    pub params: usize,
127    pub terminator: R,
128}
129impl<T, Y, R, D> Fun<T, Y, R, D> {
130    pub fn param(&mut self, a: usize) -> Value<T, Y, R, D> {
131        return Value {
132            id: self.values.alloc(ValueDef::Param(a)),
133            idx: None,
134        };
135    }
136}
137impl<T, Y, R: Default, D> Default for Fun<T, Y, R, D> {
138    fn default() -> Self {
139        Self {
140            values: Default::default(),
141            terminator: Default::default(),
142            params: Default::default(),
143        }
144    }
145}
146impl<T: Clone, Y: Clone, R: Clone, D> Clone for Fun<T, Y, R, D> {
147    fn clone(&self) -> Self {
148        Self {
149            values: self.values.clone(),
150            params: self.params.clone(),
151            terminator: self.terminator.clone(),
152        }
153    }
154}
155
156pub struct Module<T, Y, R, D> {
157    pub code: Arena<Fun<T, Y, R, D>>,
158    pub data: Arena<D>,
159}
160
161pub struct ModuleTransform<T, Y, R, D, S, L, E, P> {
162    pub input: Module<T, Y, R, D>,
163    pub out: Module<S, L, E, P>,
164    pub datum_cache: BTreeMap<Id<D>, Id<P>>,
165    pub code_cache: BTreeMap<Id<Fun<T, Y, R, D>>, Id<Fun<S, L, E, P>>>,
166}
167pub struct FuncTransformCtx<T, Y, R, D, S, L, E, P> {
168    pub input: Id<Fun<T, Y, R, D>>,
169    pub output: Id<Fun<S, L, E, P>>,
170}
171pub trait MTBehavior<T, Y, R, D, S, L, E, P> {
172    fn value(
173        &mut self,
174        ctx: &mut ModuleTransform<T, Y, R, D, S, L, E, P>,
175        fun_ctx: FuncTransformCtx<T, Y, R, D, S, L, E, P>,
176        def: ValueDef<T, Y, R, D>,
177        value: impl FnMut(
178            &mut Self,
179            &mut ModuleTransform<T, Y, R, D, S, L, E, P>,
180            Id<ValueDef<T, Y, R, D>>,
181        ) -> Id<ValueDef<S, L, E, P>>,
182        func: impl FnMut(
183            &mut Self,
184            &mut ModuleTransform<T, Y, R, D, S, L, E, P>,
185            Id<Fun<T, Y, R, D>>,
186        ) -> Id<Fun<S, L, E, P>>,
187        dat: impl FnMut(&mut Self, &mut ModuleTransform<T, Y, R, D, S, L, E, P>, Id<D>) -> Id<P>,
188    ) -> ValueDef<S, L, E, P>;
189    fn terminator(
190        &mut self,
191        ctx: &mut ModuleTransform<T, Y, R, D, S, L, E, P>,
192        fun_ctx: FuncTransformCtx<T, Y, R, D, S, L, E, P>,
193        def: R,
194        value: impl FnMut(
195            &mut Self,
196            &mut ModuleTransform<T, Y, R, D, S, L, E, P>,
197            Id<ValueDef<T, Y, R, D>>,
198        ) -> Id<ValueDef<S, L, E, P>>,
199        func: impl FnMut(
200            &mut Self,
201            &mut ModuleTransform<T, Y, R, D, S, L, E, P>,
202            Id<Fun<T, Y, R, D>>,
203        ) -> Id<Fun<S, L, E, P>>,
204        dat: impl FnMut(&mut Self, &mut ModuleTransform<T, Y, R, D, S, L, E, P>, Id<D>) -> Id<P>,
205    ) -> E;
206    fn datum(
207        &mut self,
208        ctx: &mut ModuleTransform<T, Y, R, D, S, L, E, P>,
209        def: D,
210        func: impl FnMut(
211            &mut Self,
212            &mut ModuleTransform<T, Y, R, D, S, L, E, P>,
213            Id<Fun<T, Y, R, D>>,
214        ) -> Id<Fun<S, L, E, P>>,
215        dat: impl FnMut(&mut Self, &mut ModuleTransform<T, Y, R, D, S, L, E, P>, Id<D>) -> Id<P>,
216    ) -> P;
217}
218pub type ValueTransMap<T, Y, R, D, S, L, E, P> =
219    Rc<RefCell<BTreeMap<Id<ValueDef<T, Y, R, D>>, Id<ValueDef<S, L, E, P>>>>>;
220impl<T: Clone, Y: Clone, R: Clone, D: Clone, S, L, E: Default, P>
221    ModuleTransform<T, Y, R, D, S, L, E, P>
222{
223    pub fn func_value<W: MTBehavior<T, Y, R, D, S, L, E, P>>(
224        &mut self,
225        w: &mut W,
226        f: Id<Fun<T, Y, R, D>>,
227        m: ValueTransMap<T, Y, R, D, S, L, E, P>,
228        a: Id<ValueDef<T, Y, R, D>>,
229    ) -> Id<ValueDef<S, L, E, P>> {
230        {
231            if let Some(x) = m.borrow_mut().get(&a) {
232                return *x;
233            }
234        }
235        let i = self.out.code[*self.code_cache.get(&f).unwrap()]
236            .values
237            .alloc(ValueDef::Param(usize::MAX));
238        let v = {
239            {
240                m.borrow_mut().insert(a.clone(), i.clone())
241            };
242            w.value(
243                self,
244                FuncTransformCtx {
245                    input: f.clone(),
246                    output: *self.code_cache.get(&f).unwrap(),
247                },
248                self.input.code[f].values[a].clone(),
249                |w, t, v| t.func_value(w, f, m.clone(), v),
250                |w, t, f| t.func(w, f),
251                |w, t, d| t.dat(w, d),
252            )
253        };
254        self.out.code[*self.code_cache.get(&f).unwrap()].values[i] = v;
255        return i;
256    }
257    pub fn func<W: MTBehavior<T, Y, R, D, S, L, E, P>>(
258        &mut self,
259        w: &mut W,
260        f: Id<Fun<T, Y, R, D>>,
261    ) -> Id<Fun<S, L, E, P>> {
262        if let Some(i) = self.code_cache.get(&f) {
263            return *i;
264        }
265        let me = self.out.code.alloc(Default::default());
266        let m = Rc::new(RefCell::new(BTreeMap::new()));
267        // return self.out.code.alloc_with_id(|me| {
268        self.code_cache.insert(f.clone(), me.clone());
269        let a = &self.input.code[f.clone()];
270        let t = w.terminator(
271            self,
272            FuncTransformCtx {
273                input: f.clone(),
274                output: me.clone(),
275            },
276            self.input.code[f].terminator.clone(),
277            |w, t, v| t.func_value(w, f, m.clone(), v),
278            |w, t, f| t.func(w, f),
279            |w, t, d| t.dat(w, d),
280        );
281        self.out.code[me].terminator = t;
282        return me;
283        // });
284    }
285    pub fn dat<W: MTBehavior<T, Y, R, D, S, L, E, P>>(&mut self, w: &mut W, dd: Id<D>) -> Id<P> {
286        if let Some(b) = self.datum_cache.get(&dd) {
287            return *b;
288        }
289        let d = w.datum(
290            self,
291            self.input.data[dd].clone(),
292            |w, t, f| t.func(w, f),
293            |w, t, d| t.dat(w, d),
294        );
295        let e = self.out.data.alloc(d);
296        self.datum_cache.insert(dd, e);
297        return e;
298    }
299}
300#[cfg(test)]
301mod tests {
302    use super::*;
303}
portal_ir/lib.rs

portal_ir/
lib.rs
