py_ir/
lib.rs

1pub mod types;
2pub mod value;
3
4pub trait IRValue {
5    type AssignValue: serde::Serialize + for<'a> serde::Deserialize<'a> + std::fmt::Debug + Clone;
6    type VarDefineType: serde::Serialize + for<'a> serde::Deserialize<'a> + std::fmt::Debug + Clone;
7    type FnDefineType: serde::Serialize + for<'a> serde::Deserialize<'a> + std::fmt::Debug + Clone;
8    type ParameterType: serde::Serialize + for<'a> serde::Deserialize<'a> + std::fmt::Debug + Clone;
9}
10
11#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
12pub enum Item<Var: IRValue = crate::value::Value> {
13    FnDefine(FnDefine<Var>),
14}
15
16#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
17pub struct Statements<Var: IRValue> {
18    pub stmts: Vec<Statement<Var>>,
19    pub returned: bool,
20}
21
22impl<Var: IRValue> Statements<Var> {
23    pub fn new() -> Self {
24        Self {
25            stmts: vec![],
26            returned: false,
27        }
28    }
29
30    pub fn push(&mut self, stmt: impl Into<Statement<Var>>) {
31        let stmt = stmt.into();
32        self.returned = self.returned || stmt.returned();
33        self.stmts.push(stmt);
34    }
35}
36
37impl<Var: IRValue> Default for Statements<Var> {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl<Var: IRValue> From<Vec<Statement<Var>>> for Statements<Var> {
44    fn from(stmts: Vec<Statement<Var>>) -> Self {
45        Statements {
46            returned: stmts.iter().any(|stmt| stmt.returned()),
47            stmts,
48        }
49    }
50}
51
52/// just make code generation easier
53impl<Var: IRValue> std::ops::Deref for Statements<Var> {
54    type Target = Vec<Statement<Var>>;
55
56    fn deref(&self) -> &Self::Target {
57        &self.stmts
58    }
59}
60
61/// rename vardefine support
62impl<Var: IRValue> std::ops::DerefMut for Statements<Var> {
63    fn deref_mut(&mut self) -> &mut Self::Target {
64        &mut self.stmts
65    }
66}
67
68pub trait ControlFlow {
69    fn returned(&self) -> bool;
70}
71
72impl<Var: IRValue> ControlFlow for Statements<Var> {
73    fn returned(&self) -> bool {
74        self.returned
75    }
76}
77
78#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
79pub enum Statement<Var: IRValue> {
80    VarDefine(VarDefine<Var>),
81    VarStore(VarStore<Var>),
82    Block(Statements<Var>),
83    If(If<Var>),
84    While(While<Var>),
85    Return(Return<Var>),
86}
87
88impl<Var: IRValue> ControlFlow for Statement<Var> {
89    fn returned(&self) -> bool {
90        match self {
91            Statement::Block(v) => v.returned(),
92            Statement::If(v) => v.returned(),
93            Statement::While(v) => v.returned(),
94            Statement::Return(v) => v.returned(),
95            _ => false,
96        }
97    }
98}
99
100mod from_impls {
101    use super::*;
102    impl<Var: IRValue> From<FnDefine<Var>> for Item<Var> {
103        fn from(v: FnDefine<Var>) -> Self {
104            Self::FnDefine(v)
105        }
106    }
107
108    impl<Var: IRValue> From<VarDefine<Var>> for Statement<Var> {
109        fn from(v: VarDefine<Var>) -> Self {
110            Self::VarDefine(v)
111        }
112    }
113
114    impl<Var: IRValue> From<VarStore<Var>> for Statement<Var> {
115        fn from(v: VarStore<Var>) -> Self {
116            Self::VarStore(v)
117        }
118    }
119
120    impl<Var: IRValue> From<Statements<Var>> for Statement<Var> {
121        fn from(v: Statements<Var>) -> Self {
122            Self::Block(v)
123        }
124    }
125
126    impl<Var: IRValue> From<If<Var>> for Statement<Var> {
127        fn from(v: If<Var>) -> Self {
128            Self::If(v)
129        }
130    }
131
132    impl<Var: IRValue> From<While<Var>> for Statement<Var> {
133        fn from(v: While<Var>) -> Self {
134            Self::While(v)
135        }
136    }
137
138    impl<Var: IRValue> From<Return<Var>> for Statement<Var> {
139        fn from(v: Return<Var>) -> Self {
140            Self::Return(v)
141        }
142    }
143}
144
145#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
146pub struct Parameter<Pty> {
147    #[serde(rename = "type")]
148    pub ty: Pty,
149    /// using string because its the name of parameter, not a value
150    pub name: String,
151}
152
153#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
154pub struct FnDefine<Var: IRValue> {
155    pub export: bool,
156    #[serde(rename = "type")]
157    pub ty: Var::FnDefineType,
158    pub name: String,
159    pub params: Vec<Parameter<Var::ParameterType>>,
160    pub body: Statements<Var>,
161}
162
163#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
164pub struct VarDefine<Var: IRValue> {
165    /// the type of the variable
166    #[serde(rename = "type")]
167    pub ty: Var::VarDefineType,
168    /// the name of the variable, either a named variable or a temp variable from computing
169    pub name: String,
170    /// the initial value of the variable
171    ///
172    /// [`None`] represents the variable is not initialized when it was created
173    pub init: Option<Var::AssignValue>,
174    /// a `temp value` will only be immediately used once in expressions
175    ///
176    /// this is a hint for code generation backend optimization
177    ///
178    /// code generation backend can directly inline `init` only when `is_temp` is true
179    pub is_temp: bool,
180}
181
182#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
183pub struct VarStore<Var> {
184    pub name: String,
185    pub val: Var,
186}
187
188#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
189pub struct Condition<Var: IRValue> {
190    // the final value of the condition
191    pub val: Var,
192    pub compute: Statements<Var>,
193}
194
195#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
196pub struct IfBranch<Var: IRValue> {
197    pub cond: Condition<Var>,
198    pub body: Statements<Var>,
199}
200
201impl<Var: IRValue> ControlFlow for IfBranch<Var> {
202    fn returned(&self) -> bool {
203        self.body.returned()
204    }
205}
206
207#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
208pub struct If<Var: IRValue> {
209    pub branches: Vec<IfBranch<Var>>,
210    #[serde(rename = "else")]
211    pub else_: Option<Statements<Var>>,
212}
213
214impl<Var: IRValue> ControlFlow for If<Var> {
215    fn returned(&self) -> bool {
216        self.else_.as_ref().is_some_and(|else_| else_.returned())
217            && self.branches.iter().all(|branch| branch.returned())
218    }
219}
220
221#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
222pub struct While<Var: IRValue> {
223    pub cond: Condition<Var>,
224    pub body: Statements<Var>,
225}
226
227impl<Var: IRValue> ControlFlow for While<Var> {
228    fn returned(&self) -> bool {
229        false
230    }
231}
232
233#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
234pub struct Return<Var> {
235    pub val: Option<Var>,
236}
237
238impl<Var: IRValue> ControlFlow for Return<Var> {
239    fn returned(&self) -> bool {
240        true
241    }
242}
243
244#[macro_export]
245macro_rules! custom_ir_variable {
246    ($vis:vis IR<$variable:ty>) => {
247        $vis type Item        = $crate::Item        <$variable>;
248        $vis type FnDefine    = $crate::FnDefine    <$variable>;
249        $vis type Statements  = $crate::Statements  <$variable>;
250        $vis type Statement   = $crate::Statement   <$variable>;
251        $vis type VarDefine   = $crate::VarDefine   <$variable>;
252        $vis type VarStore    = $crate::VarStore    <$variable>;
253        $vis type Condition   = $crate::Condition   <$variable>;
254        $vis type If          = $crate::If          <$variable>;
255        $vis type IfBranch    = $crate::IfBranch    <$variable>;
256        $vis type While       = $crate::While       <$variable>;
257        $vis type Return      = $crate::Return      <$variable>;
258        $vis type Parameter   = $crate::Parameter   <<$variable as $crate::IRValue>::ParameterType>;
259    };
260}