  1
-  2
-  3
-  4
-  5
-  6
-  7
-  8
-  9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
- 33
- 34
- 35
- 36
- 37
- 38
- 39
- 40
- 41
- 42
- 43
- 44
- 45
- 46
- 47
- 48
- 49
- 50
- 51
- 52
- 53
- 54
- 55
- 56
- 57
- 58
- 59
- 60
- 61
- 62
- 63
- 64
- 65
- 66
- 67
- 68
- 69
- 70
- 71
- 72
- 73
- 74
- 75
- 76
- 77
- 78
- 79
- 80
- 81
- 82
- 83
- 84
- 85
- 86
- 87
- 88
- 89
- 90
- 91
- 92
- 93
- 94
- 95
- 96
- 97
- 98
- 99
+codegen.rs - source
+        
+    
+    
+        

+                                
+                        
1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+11
+12
+13
+14
+15
+16
+17
+18
+19
+20
+21
+22
+23
+24
+25
+26
+27
+28
+29
+30
+31
+32
+33
+34
+35
+36
+37
+38
+39
+40
+41
+42
+43
+44
+45
+46
+47
+48
+49
+50
+51
+52
+53
+54
+55
+56
+57
+58
+59
+60
+61
+62
+63
+64
+65
+66
+67
+68
+69
+70
+71
+72
+73
+74
+75
+76
+77
+78
+79
+80
+81
+82
+83
+84
+85
+86
+87
+88
+89
+90
+91
+92
+93
+94
+95
+96
+97
+98
+99
 100
 101
 102
@@ -160,12 +166,171 @@
 160
 161
 162
-
-use std::collections::HashMap;
+163
+164
+165
+166
+167
+168
+169
+170
+171
+172
+173
+174
+175
+176
+177
+178
+179
+180
+181
+182
+183
+184
+185
+186
+187
+188
+189
+190
+191
+192
+193
+194
+195
+196
+197
+198
+199
+200
+201
+202
+203
+204
+205
+206
+207
+208
+209
+210
+211
+212
+213
+214
+215
+216
+217
+218
+219
+220
+221
+222
+223
+224
+225
+226
+227
+228
+229
+230
+231
+232
+233
+234
+235
+236
+237
+238
+239
+240
+241
+242
+243
+244
+245
+246
+247
+248
+249
+250
+251
+252
+253
+254
+255
+256
+257
+258
+259
+260
+261
+262
+263
+264
+265
+266
+267
+268
+269
+270
+271
+272
+273
+274
+275
+276
+277
+278
+279
+280
+281
+282
+283
+284
+285
+286
+287
+288
+289
+290
+291
+292
+293
+294
+295
+296
+297
+298
+299
+300
+301
+302
+303
+304
+305
+306
+307
+308
+309
+310
+311
+312
+313
+314
+315
+316
+317
+318
+319
+320
+321
+322
+
use std::collections::HashMap;
 
-use crate::llvm::{FnValue, FunctionPassManager, IRBuilder, Module, Value};
-use crate::parser::{ExprAST, FunctionAST, PrototypeAST};
-use crate::Either;
+use crate::llvm::{FnValue, FunctionPassManager, IRBuilder, Module, Value};
+use crate::parser::{ExprAST, FunctionAST, PrototypeAST};
+use crate::Either;
 
 type CodegenResult<T> = Result<T, String>;
 
@@ -181,9 +346,9 @@
     /// Compile either a [`PrototypeAST`] or a [`FunctionAST`] into the LLVM `module`.
     pub fn compile(
         module: &'llvm Module,
-        fn_protos: &mut HashMap<String, PrototypeAST>,
+        fn_protos: &mut HashMap<String, PrototypeAST>,
         compilee: Either<&PrototypeAST, &FunctionAST>,
-    ) -> CodegenResult<FnValue<'llvm>> {
+    ) -> CodegenResult<FnValue<'llvm>> {
         let mut cg = Codegen {
             module,
             builder: &IRBuilder::with_ctx(module),
@@ -193,40 +358,40 @@
         let mut variables = HashMap::new();
 
         match compilee {
-            Either::A(proto) => Ok(cg.codegen_prototype(proto)),
-            Either::B(func) => cg.codegen_function(func, &mut variables),
+            Either::A(proto) => Ok(cg.codegen_prototype(proto)),
+            Either::B(func) => cg.codegen_function(func, &mut variables),
         }
     }
 
     fn codegen_expr(
         &self,
         expr: &ExprAST,
-        named_values: &mut HashMap<&'llvm str, Value<'llvm>>,
-    ) -> CodegenResult<Value<'llvm>> {
+        named_values: &mut HashMap<String, Value<'llvm>>,
+    ) -> CodegenResult<Value<'llvm>> {
         match expr {
-            ExprAST::Number(num) => Ok(self.module.type_f64().const_f64(*num)),
-            ExprAST::Variable(name) => match named_values.get(name.as_str()) {
-                Some(value) => Ok(*value),
-                None => Err("Unknown variable name".into()),
+            ExprAST::Number(num) => Ok(self.module.type_f64().const_f64(*num)),
+            ExprAST::Variable(name) => match named_values.get(name.as_str()) {
+                Some(value) => Ok(*value),
+                None => Err("Unknown variable name".into()),
             },
-            ExprAST::Binary(binop, lhs, rhs) => {
+            ExprAST::Binary(binop, lhs, rhs) => {
                 let l = self.codegen_expr(lhs, named_values)?;
                 let r = self.codegen_expr(rhs, named_values)?;
 
                 match binop {
-                    '+' => Ok(self.builder.fadd(l, r)),
-                    '-' => Ok(self.builder.fsub(l, r)),
-                    '*' => Ok(self.builder.fmul(l, r)),
-                    '<' => {
+                    '+' => Ok(self.builder.fadd(l, r)),
+                    '-' => Ok(self.builder.fsub(l, r)),
+                    '*' => Ok(self.builder.fmul(l, r)),
+                    '<' => {
                         let res = self.builder.fcmpult(l, r);
                         // Turn bool into f64.
                         Ok(self.builder.uitofp(res, self.module.type_f64()))
                     }
-                    _ => Err("invalid binary operator".into()),
+                    _ => Err("invalid binary operator".into()),
                 }
             }
-            ExprAST::Call(callee, args) => match self.get_function(callee) {
-                Some(callee) => {
+            ExprAST::Call(callee, args) => match self.get_function(callee) {
+                Some(callee) => {
                     if callee.args() != args.len() {
                         return Err("Incorrect # arguments passed".into());
                     }
@@ -237,21 +402,181 @@
                         .map(|arg| self.codegen_expr(arg, named_values))
                         .collect::<CodegenResult<_>>()?;
 
-                    Ok(self.builder.call(callee, &mut args))
+                    Ok(self.builder.call(callee, &mut args))
                 }
-                None => Err("Unknown function referenced".into()),
+                None => Err("Unknown function referenced".into()),
             },
+            ExprAST::If { cond, then, else_ } => {
+                // For 'if' expressions we are building the following CFG.
+                //
+                //         ; cond
+                //         br
+                //          |
+                //    +-----+------+
+                //    v            v
+                //  ; then       ; else
+                //    |            |
+                //    +-----+------+
+                //          v
+                //        ; merge
+                //        phi then, else
+                //        ret phi
+
+                let cond_v = {
+                    // Codgen 'cond' expression.
+                    let v = self.codegen_expr(cond, named_values)?;
+                    // Compare 'v' against '0' as 'one = ordered not equal'.
+                    self.builder
+                        .fcmpone(v, self.module.type_f64().const_f64(0f64))
+                };
+
+                // Get the function we are currently inserting into.
+                let the_function = self.builder.get_insert_block().get_parent();
+
+                // Create basic blocks for the 'then' / 'else' expressions as well as the return
+                // instruction ('merge').
+                //
+                // Append the 'then' basic block to the function, don't insert the 'else' and
+                // 'merge' basic blocks yet.
+                let then_bb = self.module.append_basic_block(the_function);
+                let else_bb = self.module.create_basic_block();
+                let merge_bb = self.module.create_basic_block();
+
+                // Create a conditional branch based on the result of the 'cond' expression.
+                self.builder.cond_br(cond_v, then_bb, else_bb);
+
+                // Move to 'then' basic block and codgen the 'then' expression.
+                self.builder.pos_at_end(then_bb);
+                let then_v = self.codegen_expr(then, named_values)?;
+                // Create unconditional branch to 'merge' block.
+                self.builder.br(merge_bb);
+                // Update reference to current basic block (in case the 'then' expression added new
+                // basic blocks).
+                let then_bb = self.builder.get_insert_block();
+
+                // Now append the 'else' basic block to the function.
+                the_function.append_basic_block(else_bb);
+                // Move to 'else' basic block and codgen the 'else' expression.
+                self.builder.pos_at_end(else_bb);
+                let else_v = self.codegen_expr(else_, named_values)?;
+                // Create unconditional branch to 'merge' block.
+                self.builder.br(merge_bb);
+                // Update reference to current basic block (in case the 'else' expression added new
+                // basic blocks).
+                let else_bb = self.builder.get_insert_block();
+
+                // Now append the 'merge' basic block to the function.
+                the_function.append_basic_block(merge_bb);
+                // Move to 'merge' basic block.
+                self.builder.pos_at_end(merge_bb);
+                // Codegen the phi node returning the appropriate value depending on the branch
+                // condition.
+                let phi = self.builder.phi(
+                    self.module.type_f64(),
+                    &[(then_v, then_bb), (else_v, else_bb)],
+                );
+
+                Ok(*phi)
+            }
+            ExprAST::For {
+                var,
+                start,
+                end,
+                step,
+                body,
+            } => {
+                // For 'for' expression we build the following structure.
+                //
+                // entry:
+                //   init = start expression
+                //   br loop
+                // loop:
+                //   i = phi [%init, %entry], [%new_i, %loop]
+                //   ; loop body ...
+                //   new_i = increment %i by step expression
+                //   ; check end condition and branch
+                // end:
+
+                // Compute initial value for the loop variable.
+                let start_val = self.codegen_expr(start, named_values)?;
+
+                let the_function = self.builder.get_insert_block().get_parent();
+                // Get current basic block (used in the loop variable phi node).
+                let entry_bb = self.builder.get_insert_block();
+                // Add new basic block to emit loop body.
+                let loop_bb = self.module.append_basic_block(the_function);
+
+                self.builder.br(loop_bb);
+                self.builder.pos_at_end(loop_bb);
+
+                // Build phi not to pick loop variable in case we come from the 'entry' block.
+                // Which is the case when we enter the loop for the first time.
+                // We will add another incoming value once we computed the updated loop variable
+                // below.
+                let variable = self
+                    .builder
+                    .phi(self.module.type_f64(), &[(start_val, entry_bb)]);
+
+                // Insert the loop variable into the named values map that it can be referenced
+                // from the body as well as the end condition.
+                // In case the loop variable shadows an existing variable remember the shared one.
+                let old_val = named_values.insert(var.into(), *variable);
+
+                // Generate the loop body.
+                self.codegen_expr(body, named_values)?;
+
+                // Generate step value expression if available else use '1'.
+                let step_val = if let Some(step) = step {
+                    self.codegen_expr(step, named_values)?
+                } else {
+                    self.module.type_f64().const_f64(1f64)
+                };
+
+                // Increment loop variable.
+                let next_var = self.builder.fadd(*variable, step_val);
+
+                // Generate the loop end condition.
+                let end_cond = self.codegen_expr(end, named_values)?;
+                let end_cond = self
+                    .builder
+                    .fcmpone(end_cond, self.module.type_f64().const_f64(0f64));
+
+                // Get current basic block.
+                let loop_end_bb = self.builder.get_insert_block();
+                // Add new basic block following the loop.
+                let after_bb = self.module.append_basic_block(the_function);
+
+                // Register additional incoming value for the loop variable. This will choose the
+                // updated loop variable if we are iterating in the loop.
+                variable.add_incoming(next_var, loop_end_bb);
+
+                // Branch depending on the loop end condition.
+                self.builder.cond_br(end_cond, loop_bb, after_bb);
+
+                self.builder.pos_at_end(after_bb);
+
+                // Restore the shadowed variable if there was one.
+                if let Some(old_val) = old_val {
+                    // We inserted 'var' above so it must exist.
+                    *named_values.get_mut(var).unwrap() = old_val;
+                } else {
+                    named_values.remove(var);
+                }
+
+                // Loops just always return 0.
+                Ok(self.module.type_f64().const_f64(0f64))
+            }
         }
     }
 
-    fn codegen_prototype(&self, PrototypeAST(name, args): &PrototypeAST) -> FnValue<'llvm> {
+    fn codegen_prototype(&self, PrototypeAST(name, args): &PrototypeAST) -> FnValue<'llvm> {
         let type_f64 = self.module.type_f64();
 
         let mut doubles = Vec::new();
         doubles.resize(args.len(), type_f64);
 
         // Build the function type: fn(f64, f64, ..) -> f64
-        let ft = self.module.type_fn(&mut doubles, type_f64);
+        let ft = self.module.type_fn(&mut doubles, type_f64);
 
         // Create the function declaration.
         let f = self.module.add_fn(name, ft);
@@ -265,10 +590,10 @@
     }
 
     fn codegen_function(
-        &mut self,
+        &mut self,
         FunctionAST(proto, body): &FunctionAST,
-        named_values: &mut HashMap<&'llvm str, Value<'llvm>>,
-    ) -> CodegenResult<FnValue<'llvm>> {
+        named_values: &mut HashMap<String, Value<'llvm>>,
+    ) -> CodegenResult<FnValue<'llvm>> {
         // Insert the function prototype into the `fn_protos` map to keep track for re-generating
         // declarations in other modules.
         self.fn_protos.insert(proto.0.clone(), proto.clone());
@@ -290,7 +615,7 @@
         // Update the map with the current functions args.
         for idx in 0..the_function.args() {
             let arg = the_function.arg(idx);
-            named_values.insert(arg.get_name(), arg);
+            named_values.insert(arg.get_name().into(), arg);
         }
 
         // Codegen function body.
@@ -311,10 +636,10 @@
     /// If the function is not available in the module, check if the prototype is known and codegen
     /// it.
     /// Return [`None`] if the prototype is not known.
-    fn get_function(&self, name: &str) -> Option<FnValue<'llvm>> {
+    fn get_function(&self, name: &str) -> Option<FnValue<'llvm>> {
         let callee = match self.module.get_fn(name) {
-            Some(callee) => callee,
-            None => {
+            Some(callee) => callee,
+            None => {
                 let proto = self.fn_protos.get(name)?;
                 self.codegen_prototype(proto)
             }
@@ -323,7 +648,6 @@
         Some(callee)
     }
 }
-

-

-    
+
+