1use std::iter::empty;
2
3use arena_traits::{Arena, IndexIter};
4use either::Either;
5use proc_macro2::{Span, TokenStream};
6use quasiquote::quasiquote;
7use quote::{format_ident, quote};
8use relooper::{BranchMode, RelooperLabel, ShapedBlock};
9use ssa_traits::{Block, Target, Term, TypedBlock, TypedFunc, TypedValue};
10use cfg_traits::{Block as CFGBlock, Target as CFGTarget, Term as CFGTerm};
11use syn::{Ident, Lifetime};
12fn term(b: &BranchMode) -> TokenStream {
13    match b {
14        relooper::BranchMode::LoopBreak(l) => {
15            let l = Lifetime::new(&format!("'l{}", l), Span::call_site());
16            quote! {
17                break #l;
18            }
19        }
20        relooper::BranchMode::LoopBreakIntoMulti(l) => {
21            let l = Lifetime::new(&format!("'l{}", l), Span::call_site());
22            quote! {
23                break #l;
24            }
25        }
26        relooper::BranchMode::LoopContinue(l) => {
27            let l = Lifetime::new(&format!("'l{}", l), Span::call_site());
28            quote! {
29                continue #l;
30            }
31        }
32        relooper::BranchMode::LoopContinueIntoMulti(l) => {
33            let l = Lifetime::new(&format!("'l{}", l), Span::call_site());
34            quote! {
35                continue #l;
36            }
37        }
38        relooper::BranchMode::MergedBranch => {
39            quote! {}
40        }
41        relooper::BranchMode::MergedBranchIntoMulti => quote! {},
42        relooper::BranchMode::SetLabelAndBreak => quote! {
43            break 'cff;
44        },
45    }
46}
47pub trait RsFunc:
48    TypedFunc<
49        Ty: Rs<Self>,
50        Block: RsId<Self> + Ord + Clone + RelooperLabel,
51        Value: RsId<Self> + Clone,
52        Values: Arena<Self::Value, Output: Rs<Self>>,
53        Blocks: Arena<Self::Block, Output: Block<Self, Terminator: RsTerm<Self>>>,
54    > + Sized
55{
56}
57impl<
58        T: TypedFunc<
59                Ty: Rs<Self>,
60                Block: RsId<Self> + Ord + Clone + RelooperLabel,
61                Value: RsId<Self> + Clone,
62                Values: Arena<Self::Value, Output: Rs<Self>>,
63                Blocks: Arena<Self::Block, Output: Block<Self, Terminator: RsTerm<Self>>>,
64            > + Sized,
65    > RsFunc for T
66{
67}
68pub trait Rs<F: RsFunc> {
69    fn rs(&self, f: &F) -> anyhow::Result<TokenStream>;
70}
71pub trait RsId<F: RsFunc> {
72    fn rs(&self, f: &F) -> anyhow::Result<Ident>;
73}
74pub trait RsTerm<F: RsFunc> {
75    fn rs_term(
76        &self,
77        f: &F,
78        go: impl FnMut(F::Block) -> anyhow::Result<TokenStream>,
79    ) -> anyhow::Result<TokenStream>;
80}
81pub fn render_target<R: RsFunc>(
82    t: &impl Target<R>,
83    f: &R,
84    go: &mut impl FnMut(R::Block) -> anyhow::Result<TokenStream>,
85    prepend: impl Iterator<Item: Rs<R>>,
86) -> anyhow::Result<TokenStream> {
87    let vars = prepend
88        .map(|x| x.rs(f))
89        .chain(t.values(f).map(|a| {
90            let a = format_ident!("V{}", a.rs(f)?);
91            anyhow::Ok(quasiquote!(#a .take()))
92        }))
93        .enumerate()
94        .map(|(i, a)| {
95            let i = format_ident!("P{}_{i}", t.block().rs(f)?);
96            Ok(quasiquote! {
97                #i = #{a?};
98            })
99        })
100        .collect::<anyhow::Result<Vec<_>>>()?;
101    Ok(quasiquote! {
102        #(#vars);*
103        #{go(t.block())?}
104    })
105}
106pub fn go<F: RsFunc>(params: &[TokenStream], f: &F, e: F::Block) -> anyhow::Result<TokenStream> {
107    Ok(quasiquote! {
108        #{
109            let k = f.blocks().iter().flat_map(|a|{
110                if a == e{
111                    Either::Left(params.iter().enumerate().map(move|(i,p)|Ok(quasiquote!(let mut #{format_ident!("P{}_{i}",a.rs(f)?)} = Some(#p)))))
112                }else{
113                    Either::Right(f.blocks()[a].params().map(|(a,_)|a).enumerate().map(move|(i,t)|Ok(quasiquote!(let mut #{format_ident!("P{}_{i}",a.rs(f)?)}: Option<#{t.rs(f)?}> = None))))
114                }
115            }).chain(f.values().iter().map(|v|{
116                let ty = f.values()[v.clone()].ty(f).rs(f)?;
117                Ok(quote!{
118                    let mut #{format_ident!("V{}",v.rs(f)?)}: Option<#ty> = None;
119                })
120            })).collect::<anyhow::Result<Vec<_>>>()?;
121            quote! {
122                #(#k);*
123            }
124        };
125        let mut cff = 0usize;
126        #{block(f,ssa_reloop::go(f,e).as_ref())?}
127    })
128}
129pub fn idx_of<F: RsFunc>(f: &F, k: F::Block) -> usize {
130    f.blocks()
131        .iter()
132        .enumerate()
133        .find_map(|(i, l)| if l == k { Some(i) } else { None })
134        .unwrap()
135}
136pub fn block<F: RsFunc>(f: &F, k: &ShapedBlock<F::Block>) -> anyhow::Result<TokenStream> {
137    match k {
138        ShapedBlock::Loop(l) => {
139            let r = block(f, &l.inner.as_ref())?;
140            let next = l.next.as_ref();
141            let next = match next {
142                None => Default::default(),
143                Some(a) => block(f, a)?,
144            };
145            let l = Lifetime::new(&format!("'l{}", l.loop_id), Span::call_site());
146            Ok(quote! {
147                #l : loop{
148                    #r
149                };
150                #next;
151            })
152        }
153        ShapedBlock::Multiple(k) => {
154            let initial = k.handled.iter().enumerate().flat_map(|(a, b)| {
155                b.labels.iter().map(move |l| {
156                    let l = idx_of(f, *l);
157                    quote! {
158                        #l => #a
159                    }
160                })
161            });
162            let cases = k
163                .handled
164                .iter()
165                .enumerate()
166                .map(|(a, i)| {
167                    let ib = block(f, &i.inner)?;
168                    let ic = if i.break_after {
169                        quote! {}
170                    } else {
171                        quote! {
172                            cff2 += 1;
173                            continue 'cff
174                        }
175                    };
176                    Ok(quote! {
177                        #a => {
178                            #ib;
179                            #ic;
180                        }
181                    })
182                })
183                .collect::<anyhow::Result<Vec<_>>>()?;
184            Ok(quote! {
185                let mut cff2 = match cff{
186                    #(#initial),*,
187                    _ => unreachable!()
188                };
189                'cff: loop{
190                    match cff2{
191                        #(#cases),*,
192                        _ => unreachable!()
193                    };
194                    break 'cff;
195                };
196            })
197        }
198        ShapedBlock::Simple(s) => {
199            let immediate = match s.immediate.as_ref() {
200                None => Default::default(),
201                Some(a) => block(f, a)?,
202            };
203
204            let next = match s.next.as_ref() {
205                None => Default::default(),
206                Some(a) => block(f, a)?,
207            };
208            let stmts = f.blocks()[s.label]
209                .insts()
210                .map(|v| {
211                    Ok(quasiquote! {
212                        #{format_ident!("V{}",v.rs(f)?)} = Some(#{f.values()[v].rs(f)?})
213                    })
214                })
215                .collect::<anyhow::Result<Vec<_>>>()?;
216            let term = f.blocks()[s.label].term().rs_term(f, |k| {
217                let br = term(
218                    &s.branches
219                        .get(&k)
220                        .cloned()
221                        .unwrap_or(relooper::BranchMode::MergedBranch),
222                );
223                let bi = idx_of(f, k);
224                Ok(quote! {
225                    cff = #bi;
226                    #br
227                })
228            })?;
229            Ok(quote! {
230                #(#stmts);*;
231                #term;
232                #immediate;
233                #next;
234            })
235        }
236    }
237}
238#[cfg(feature = "id-arena")]
239impl<F: RsFunc, T> RsId<F> for id_arena::Id<T> {
240    fn rs(&self, f: &F) -> anyhow::Result<Ident> {
241        Ok(format_ident!("V{}", self.index()))
242    }
243}
244pub trait RsOp<F: RsFunc> {
245    fn rs_op(
246        &self,
247        f: &F,
248        all: &[impl RsId<F>],
249        blargs: &[F::Block],
250    ) -> anyhow::Result<TokenStream>;
251}
252#[cfg(feature = "ssa-canon")]
253impl<F: RsFunc<Block = id_arena::Id<ssa_canon::Block<O, T, Y>>>, O: RsOp<F>, T, Y> Rs<F>
254    for ssa_canon::Value<O, T, Y>
255{
256    fn rs(&self, f: &F) -> anyhow::Result<TokenStream> {
257        match self {
258            ssa_canon::Value::Op(o, args, q, _) => o.rs_op(f, &args, q.as_slice()),
259            ssa_canon::Value::Param(i, a, _) => {
260                Ok(quasiquote!(#{format_ident!("P{}_{i}",a.rs(f)?)}.take().unwrap()))
261            }
262        }
263    }
264}
265#[cfg(feature = "ssa-canon")]
266impl<
267        O: RsOp<ssa_canon::Func<O, T, Y>>,
268        T: Term<ssa_canon::Func<O, T, Y>, Target = ssa_canon::Target<O, T, Y>>,
269        Y: Clone,
270    > RsTerm<ssa_canon::Func<O, T, Y>> for ssa_canon::Target<O, T, Y>
271where
272    ssa_canon::Func<O, T, Y>: RsFunc<Block = id_arena::Id<ssa_canon::Block<O, T, Y>>>,
273{
274    fn rs_term(
275        &self,
276        f: &ssa_canon::Func<O, T, Y>,
277        mut go: impl FnMut(
278            <ssa_canon::Func<O, T, Y> as cfg_traits::Func>::Block,
279        ) -> anyhow::Result<TokenStream>,
280    ) -> anyhow::Result<TokenStream> {
281        render_target(self, f, &mut go, empty::<ssa_canon::Value<O, T, Y>>())
282    }
283}
284impl<F: RsFunc, A: RsOp<F>, B: RsOp<F>> RsOp<F> for Either<A, B> {
285    fn rs_op(
286        &self,
287        f: &F,
288        all: &[impl RsId<F>],
289        blargs: &[F::Block],
290    ) -> anyhow::Result<TokenStream> {
291        match self {
292            Either::Left(a) => a.rs_op(f, all, blargs),
293            Either::Right(b) => b.rs_op(f, all, blargs),
294        }
295    }
296}