swc_ssa/
lib.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    convert::Infallible,
4    default,
5    iter::{empty, once},
6    mem::take,
7};
8
9use anyhow::Context;
10use cfg_traits::Term;
11use id_arena::{Arena, Id};
12use portal_jsc_common::LId;
13use ssa_traits::Value;
14use swc_atoms::Atom;
15use swc_cfg::Catch;
16use swc_common::Span;
17use swc_ecma_ast::{Id as Ident, Lit, Null, TsType, TsTypeAnn, TsTypeParamDecl};
18use swc_tac::{Item, TBlock, TCallee, TCfg, TFunc, TStmt, ValFlags};
19pub mod ch;
20pub mod idw;
21pub mod impls;
22pub mod rew;
23pub mod simplify;
24
25pub fn benj(swc_func: &mut SCfg) {
26    for block_index in swc_func.blocks.iter().map(|a| a.0).collect::<Vec<_>>() {
27        let mut postcedent = take(&mut swc_func.blocks[block_index].postcedent);
28        for target in postcedent.targets_mut() {
29            if target.block.index() <= block_index.index() {
30                for arg in target.args.iter_mut() {
31                    let value = swc_func.values.alloc(SValueW { value: SValue::Benc(*arg) });
32                    swc_func.blocks[block_index].stmts.push(value);
33                    *arg = value;
34                }
35            }
36        }
37        swc_func.blocks[block_index].postcedent = postcedent;
38    }
39}
40#[derive(Clone, Debug)]
41pub struct SFunc {
42    pub cfg: SCfg,
43    pub entry: Id<SBlock>,
44    pub is_generator: bool,
45    pub is_async: bool,
46    pub ts_params: Vec<Option<TsType>>,
47}
48impl TryFrom<TFunc> for SFunc {
49    type Error = anyhow::Error;
50
51    fn try_from(value: TFunc) -> Result<Self, Self::Error> {
52        let mut decls = value.cfg.decls.clone();
53        let mut d = BTreeSet::new();
54        for e in value.cfg.externs().collect::<BTreeSet<_>>() {
55            decls.remove(&e);
56            d.insert(e);
57        }
58        let mut cfg = SCfg {
59            blocks: Default::default(),
60            values: Default::default(),
61            ts: Default::default(),
62            decls: d,
63            generics: None,
64            ts_retty: None,
65        };
66        let entry2 = cfg.blocks.alloc(Default::default());
67        let params = value
68            .params
69            .iter()
70            .cloned()
71            .map(|a| (a, cfg.add_blockparam(entry2)))
72            .collect::<BTreeMap<_, _>>();
73        let undef = cfg.values.alloc(
74            SValue::Item {
75                item: Item::Undef,
76                span: None,
77            }
78            .into(),
79        );
80        let mut trans = Trans {
81            map: BTreeMap::new(),
82            all: decls.clone(),
83            undef,
84        };
85        let entry = trans.trans(&value.cfg, &mut cfg, value.entry)?;
86        let target = STarget {
87            block: entry,
88            args: decls
89                .iter()
90                .cloned()
91                .map(|a| match params.get(&a) {
92                    Some(v) => *v,
93                    None => undef,
94                })
95                .collect(),
96        };
97        cfg.blocks[entry2].postcedent.term = STerm::Jmp(target);
98        cfg.generics = value.cfg.generics;
99        cfg.ts_retty = value.cfg.ts_retty;
100        Ok(Self {
101            cfg,
102            entry: entry2,
103            is_generator: value.is_generator,
104            is_async: value.is_async,
105            ts_params: value.ts_params,
106        })
107    }
108}
109#[derive(Default, Clone, Debug)]
110pub struct SCfg {
111    pub blocks: Arena<SBlock>,
112    pub values: Arena<SValueW>,
113    pub ts: BTreeMap<Id<SValueW>, TsType>,
114    pub decls: BTreeSet<Ident>,
115    pub generics: Option<TsTypeParamDecl>,
116    pub ts_retty: Option<TsTypeAnn>,
117}
118impl SCfg {
119    pub fn refs(&self) -> BTreeSet<Ident> {
120        return self
121            .values
122            .iter()
123            .flat_map(|(a, b)| match &b.value {
124                SValue::LoadId(target) | SValue::StoreId { target, val: _ } => {
125                    [target.clone()].into_iter().collect::<BTreeSet<Ident>>()
126                }
127                SValue::Item {
128                    item: Item::Func { func },
129                    span,
130                } => func.cfg.externals(),
131                _ => Default::default(),
132            })
133            .collect();
134    }
135    pub fn externals(&self) -> BTreeSet<Ident> {
136        return self
137            .refs()
138            .into_iter()
139            .filter(|a| !self.decls.contains(&*a))
140            .collect();
141    }
142}
143#[derive(Default, Clone, Debug)]
144pub struct SBlock {
145    pub params: Vec<(Id<SValueW>, ())>,
146    pub stmts: Vec<Id<SValueW>>,
147    pub postcedent: SPostcedent,
148}
149#[derive(Clone, Debug)]
150pub struct SPostcedent<I = Id<SValueW>, B = Id<SBlock>> {
151    pub term: STerm<I, B>,
152    pub catch: SCatch<I, B>,
153}
154impl<I, B> Default for SPostcedent<I, B> {
155    fn default() -> Self {
156        Self {
157            term: Default::default(),
158            catch: Default::default(),
159        }
160    }
161}
162
163#[derive(Clone, Debug)]
164#[non_exhaustive]
165pub enum SValue<I = Id<SValueW>, B = Id<SBlock>, F = SFunc> {
166    Param {
167        block: B,
168        idx: usize,
169        ty: (),
170    },
171    Item {
172        item: Item<I, F>,
173        span: Option<Span>,
174    },
175    Assign {
176        target: LId<I>,
177        val: I,
178    },
179    LoadId(Ident),
180    StoreId {
181        target: Ident,
182        val: I,
183    },
184    Benc(I),
185}
186impl<I: Copy, B, F> SValue<I, B, F> {
187    pub fn vals<'a>(&'a self) -> Box<dyn Iterator<Item = I> + 'a> {
188        match self {
189            SValue::Param { block, idx, ty } => Box::new(empty()),
190            SValue::Item { item, span } => Box::new(item.refs().map(|a| *a)),
191            SValue::Assign { target, val } => {
192                let v = once(*val);
193                let w: Box<dyn Iterator<Item = &I> + '_> = match target {
194                    LId::Id { id } => todo!(),
195                    LId::Member { obj, mem } => Box::new([obj, &mem[0]].into_iter()),
196                    _ => todo!(),
197                };
198                Box::new(v.chain(w.cloned()))
199            }
200            SValue::LoadId(_) => Box::new(empty()),
201            SValue::StoreId { target, val } => Box::new(once(*val)),
202            SValue::Benc(a) => Box::new(once(*a)),
203        }
204    }
205}
206impl<I, B, F> SValue<I, B, F> {
207    pub fn as_ref(&self) -> SValue<&I, &B, &F> {
208        match self {
209            SValue::Param { block, idx, ty } => SValue::Param {
210                block,
211                idx: *idx,
212                ty: *ty,
213            },
214            SValue::Item { item, span } => SValue::Item {
215                item: item.as_ref(),
216                span: *span,
217            },
218            SValue::Assign { target, val } => SValue::Assign {
219                target: target
220                    .as_ref()
221                    .map2(&mut (), &mut |_, a| Ok::<_, Infallible>(a), &mut |_, b| {
222                        Ok([&b[0]])
223                    })
224                    .unwrap(),
225                val,
226            },
227            SValue::LoadId(v) => SValue::LoadId(v.clone()),
228            SValue::StoreId { target, val } => SValue::StoreId {
229                target: target.clone(),
230                val,
231            },
232            SValue::Benc(v) => SValue::Benc(v),
233        }
234    }
235    pub fn as_mut(&mut self) -> SValue<&mut I, &mut B, &mut F> {
236        match self {
237            SValue::Param { block, idx, ty } => SValue::Param {
238                block,
239                idx: *idx,
240                ty: *ty,
241            },
242            SValue::Item { item, span } => SValue::Item {
243                item: item.as_mut(),
244                span: *span,
245            },
246            SValue::Assign { target, val } => SValue::Assign {
247                target: target
248                    .as_mut()
249                    .map2(&mut (), &mut |_, a| Ok::<_, Infallible>(a), &mut |_, b| {
250                        Ok([&mut b[0]])
251                    })
252                    .unwrap(),
253                val,
254            },
255            SValue::LoadId(v) => SValue::LoadId(v.clone()),
256            SValue::StoreId { target, val } => SValue::StoreId {
257                target: target.clone(),
258                val,
259            },
260            SValue::Benc(v) => SValue::Benc(v),
261        }
262    }
263    pub fn map<J: Ord, C, G, X, E>(
264        self,
265        cx: &mut X,
266        ident: &mut (dyn FnMut(&mut X, I) -> Result<J, E> + '_),
267        block_: &mut (dyn FnMut(&mut X, B) -> Result<C, E> + '_),
268        fun: &mut (dyn FnMut(&mut X, F) -> Result<G, E> + '_),
269    ) -> Result<SValue<J, C, G>, E> {
270        Ok(match self {
271            SValue::Param { block, idx, ty } => SValue::Param {
272                block: block_(cx, block)?,
273                idx,
274                ty,
275            },
276            SValue::Item { item, span } => SValue::Item {
277                item: item.map2(cx, ident, fun)?,
278                span,
279            },
280            SValue::Assign { target, val } => SValue::Assign {
281                target: target.map(&mut |a| ident(cx, a))?,
282                val: ident(cx, val)?,
283            },
284            SValue::LoadId(a) => SValue::LoadId(a),
285            SValue::StoreId { target, val } => SValue::StoreId {
286                target,
287                val: ident(cx, val)?,
288            },
289            SValue::Benc(i) => SValue::Benc(ident(cx, i)?),
290        })
291    }
292    pub fn vals_ref<'a>(&'a self) -> Box<dyn Iterator<Item = &'a I> + 'a> {
293        match self {
294            SValue::Param { block, idx, ty } => Box::new(empty()),
295            SValue::Item { item, span } => Box::new(item.refs()),
296            SValue::Assign { target, val } => {
297                let v = once(val);
298                let w: Box<dyn Iterator<Item = &I> + '_> = match target {
299                    LId::Id { id } => todo!(),
300                    LId::Member { obj, mem } => Box::new([obj, &mem[0]].into_iter()),
301                    _ => todo!(),
302                };
303                Box::new(v.chain(w))
304            }
305            SValue::LoadId(_) => Box::new(empty()),
306            SValue::StoreId { target, val } => Box::new(once(val)),
307            SValue::Benc(a) => Box::new(once(a)),
308        }
309    }
310    pub fn vals_mut<'a>(&'a mut self) -> Box<dyn Iterator<Item = &'a mut I> + 'a> {
311        match self {
312            SValue::Param { block, idx, ty } => Box::new(empty()),
313            SValue::Item { item, span } => item.refs_mut(),
314            SValue::Assign { target, val } => {
315                let v = once(val);
316                let w: Box<dyn Iterator<Item = &mut I> + '_> = match target {
317                    LId::Id { id } => todo!(),
318                    LId::Member { obj, mem } => Box::new([obj, &mut mem[0]].into_iter()),
319                    _ => todo!(),
320                };
321                Box::new(v.chain(w))
322            }
323            SValue::LoadId(_) => Box::new(empty()),
324            SValue::StoreId { target, val } => Box::new(once(val)),
325            SValue::Benc(a) => Box::new(once(a)),
326        }
327    }
328}
329#[repr(transparent)]
330#[derive(Clone, Debug)]
331pub struct SValueW { pub value: SValue }
332impl From<SValue> for SValueW {
333    fn from(value: SValue) -> Self {
334        Self { value }
335    }
336}
337impl From<SValueW> for SValue {
338    fn from(value: SValueW) -> Self {
339        value.value
340    }
341}
342#[derive(Clone, Debug)]
343#[non_exhaustive]
344pub enum SCatch<I = Id<SValueW>, B = Id<SBlock>> {
345    Throw,
346    Just { target: STarget<I, B> },
347}
348impl<I, B> Default for SCatch<I, B> {
349    fn default() -> Self {
350        Self::Throw
351    }
352}
353#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
354pub struct STarget<I = Id<SValueW>, B = Id<SBlock>> {
355    pub block: B,
356    pub args: Vec<I>,
357}
358#[derive(Clone, Debug)]
359#[non_exhaustive]
360pub enum STerm<I = Id<SValueW>, B = Id<SBlock>> {
361    Throw(I),
362    Return(Option<I>),
363    Jmp(STarget<I, B>),
364    CondJmp {
365        cond: I,
366        if_true: STarget<I, B>,
367        if_false: STarget<I, B>,
368    },
369    Switch {
370        x: I,
371        blocks: Vec<(I, STarget<I, B>)>,
372        default: STarget<I, B>,
373    },
374    Default,
375}
376impl<I, B> Default for STerm<I, B> {
377    fn default() -> Self {
378        Self::Default
379    }
380}
381impl SCfg {
382    pub fn add_blockparam(&mut self, k: Id<SBlock>) -> Id<SValueW> {
383        let val = SValue::Param {
384            block: k,
385            idx: self.blocks[k].params.len(),
386            ty: (),
387        };
388        let val = self.values.alloc(val.into());
389        self.blocks[k].params.push((val, ()));
390        return val;
391    }
392}
393
394#[non_exhaustive]
395pub struct Trans {
396    pub map: BTreeMap<Id<TBlock>, Id<SBlock>>,
397    pub all: BTreeSet<Ident>,
398    pub undef: Id<SValueW>,
399}
400impl Trans {
401    pub fn from_undef(undef: Id<SValueW>) -> Self {
402        Self {
403            map: Default::default(),
404            all: Default::default(),
405            undef: undef,
406        }
407    }
408    pub fn apply_shim(
409        &self,
410        o: &mut SCfg,
411        state: &BTreeMap<Ident, (Id<SValueW>, ValFlags)>,
412        s: &Option<(Id<SBlock>, Vec<Ident>)>,
413        x: Id<SBlock>,
414    ) {
415        let Some((a, b)) = s else {
416            o.blocks[x].postcedent.catch = SCatch::Throw;
417            return;
418        };
419        let k = SCatch::Just {
420            target: STarget {
421                block: *a,
422                args: b
423                    .iter()
424                    .filter_map(|x| state.get(x))
425                    .map(|a| &a.0)
426                    .cloned()
427                    .collect(),
428            },
429        };
430        o.blocks[x].postcedent.catch = k;
431    }
432    pub fn load(
433        &self,
434        state: &BTreeMap<Ident, (Id<SValueW>, ValFlags)>,
435        i: &TCfg,
436        o: &mut SCfg,
437        t: Id<SBlock>,
438        a: Ident,
439        cache: &BTreeMap<Ident, Id<SValueW>>,
440    ) -> Id<SValueW> {
441        if let Some(k) = cache.get(&a) {
442            return *k;
443        }
444        let x = match state.get(&a).cloned() {
445            Some(b) => b.0,
446            None => {
447                let v = o.values.alloc(SValue::LoadId(a).into());
448                o.blocks[t].stmts.push(v);
449                return v;
450            }
451        };
452        if let Some(ty) = i.type_annotations.get(&a).cloned() {
453            o.ts.insert(x, ty);
454        };
455        x
456    }
457    pub fn trans(&mut self, i: &TCfg, o: &mut SCfg, k: Id<TBlock>) -> anyhow::Result<Id<SBlock>> {
458        loop {
459            if let Some(a) = self.map.get(&k) {
460                return Ok(*a);
461            }
462            let mut t = o.blocks.alloc(SBlock {
463                params: vec![],
464                stmts: vec![],
465                postcedent: SPostcedent::default(),
466            });
467            self.map.insert(k, t);
468            let shim: Option<(Id<SBlock>, Vec<Ident>)> = match &i.blocks[k].catch {
469                swc_tac::TCatch::Throw => None,
470                swc_tac::TCatch::Jump { pat, k } => {
471                    let a = o.blocks.alloc(SBlock {
472                        params: vec![],
473                        stmts: vec![],
474                        postcedent: SPostcedent::default(),
475                    });
476                    let mut state2 = once(pat.clone())
477                        .chain(self.all.iter().filter(|a| *a != pat).cloned())
478                        .collect::<Vec<_>>();
479                    let mut v = state2
480                        .iter()
481                        .cloned()
482                        .map(|a2| (a2, o.add_blockparam(a)))
483                        .collect::<BTreeMap<_, _>>();
484                    let p = self
485                        .all
486                        .iter()
487                        .filter_map(|x| v.get(x))
488                        .cloned()
489                        .collect::<Vec<_>>();
490                    let t = STerm::Jmp(STarget {
491                        block: self.trans(i, o, *k)?,
492                        args: p,
493                    });
494                    o.blocks[a].postcedent.term = t;
495                    Some((a, state2))
496                }
497            };
498            let mut state = self
499                .all
500                .iter()
501                .map(|a| (a.clone(), (o.add_blockparam(t), ValFlags::all())))
502                .collect::<BTreeMap<_, _>>();
503            self.apply_shim(o, &state, &shim, t);
504            let mut cache = BTreeMap::new();
505            for TStmt {
506                left: a,
507                flags,
508                right: b,
509                span: s,
510            } in i.blocks[k].stmts.iter()
511            {
512                let mut b = b.clone();
513                if let Item::Call { callee, args } = &mut b {
514                    if let TCallee::Val(v) = callee {
515                        if !i.blocks.iter().any(|k| {
516                            k.1.stmts.iter().any(|a| match &a.left {
517                                LId::Id { id } => id == v,
518                                _ => false,
519                            })
520                        }) {
521                            *callee = TCallee::Static(v.clone());
522                        }
523                    }
524                }
525                let b = b.map2::<_, _, anyhow::Error, ()>(
526                    &mut (),
527                    &mut |_, a| Ok(self.load(&state, i, o, t, a, &cache)),
528                    &mut |_, b| b.try_into(),
529                )?;
530                let b = o.values.alloc(
531                    SValue::Item {
532                        item: b,
533                        span: Some(*s),
534                    }
535                    .into(),
536                );
537                o.blocks[t].stmts.push(b);
538                let flags = match a.clone() {
539                    LId::Id { id } => match state.get_mut(&id) {
540                        Some((a, f)) => {
541                            *f &= *flags;
542                            let f = *f;
543                            *a = b;
544                            if !f.contains(ValFlags::SSA_LIKE) {
545                                let u = o.blocks.alloc(SBlock {
546                                    params: vec![],
547                                    stmts: vec![],
548                                    postcedent: Default::default(),
549                                });
550                                self.apply_shim(o, &state, &shim, u);
551                                o.blocks[t].postcedent.term = STerm::Jmp(STarget {
552                                    block: u,
553                                    args: vec![],
554                                });
555                                t = u;
556                            }
557                            Some(f)
558                        }
559                        None => {
560                            cache.insert(id.clone(), b);
561                            let c = o.values.alloc(
562                                SValue::StoreId {
563                                    target: id.clone(),
564                                    val: b,
565                                }
566                                .into(),
567                            );
568                            o.blocks[t].stmts.push(c);
569                            None
570                        }
571                    },
572                    a => {
573                        // let obj = self.load(&state, o, t, obj.clone());
574                        // let mem = self.load(&state, o, t, mem.clone());
575                        let c = a
576                            .map::<_, Infallible>(
577                                &mut |a| Ok(self.load(&state, i, o, t, a, &cache)),
578                            )
579                            .unwrap();
580                        let c = o.values.alloc(SValue::Assign { target: c, val: b }.into());
581                        o.blocks[t].stmts.push(c);
582                        None
583                    }
584                };
585            }
586            let params = self
587                .all
588                .iter()
589                .filter_map(|a| state.get(a))
590                .map(|a| &a.0)
591                .cloned()
592                .collect::<Vec<_>>();
593            let term = match &i.blocks[k].term {
594                swc_tac::TTerm::Return(ident) => match ident.as_ref() {
595                    None => STerm::Return(None),
596                    Some(a) => STerm::Return(Some(self.load(&state, i, o, t, a.clone(), &cache))),
597                },
598                swc_tac::TTerm::Throw(ident) => {
599                    STerm::Throw(self.load(&state, i, o, t, ident.clone(), &cache))
600                }
601                swc_tac::TTerm::Jmp(id) => {
602                    let id = self.trans(i, o, *id)?;
603                    STerm::Jmp(STarget {
604                        block: id,
605                        args: params,
606                    })
607                }
608                swc_tac::TTerm::CondJmp {
609                    cond,
610                    if_true,
611                    if_false,
612                } => {
613                    let if_true = self.trans(i, o, *if_true)?;
614                    let if_true = STarget {
615                        block: if_true,
616                        args: params.clone(),
617                    };
618                    let if_false = self.trans(i, o, *if_false)?;
619                    let if_false = STarget {
620                        block: if_false,
621                        args: params,
622                    };
623                    let cond = self.load(&state, i, o, t, cond.clone(), &cache);
624                    STerm::CondJmp {
625                        cond,
626                        if_true,
627                        if_false,
628                    }
629                }
630                swc_tac::TTerm::Switch { x, blocks, default } => {
631                    let x = self.load(&state, i, o, t, x.clone(), &cache);
632                    let blocks = blocks
633                        .iter()
634                        .map(|(a, b)| {
635                            let c = self.trans(i, o, *b)?;
636                            let d = self.load(&state, i, o, t, a.clone(), &cache);
637                            Ok((
638                                d,
639                                STarget {
640                                    block: c,
641                                    args: params.clone(),
642                                },
643                            ))
644                        })
645                        .collect::<anyhow::Result<Vec<_>>>()?;
646                    let default = self.trans(i, o, *default)?;
647                    let default = STarget {
648                        block: default,
649                        args: params,
650                    };
651                    STerm::Switch { x, blocks, default }
652                }
653                swc_tac::TTerm::Default => STerm::Default,
654            };
655            o.blocks[t].postcedent.term = term;
656        }
657    }
658}
swc_ssa/lib.rs

swc_ssa/
lib.rs
