minijinja/compiler/
codegen.rs

1use std::collections::BTreeMap;
2
3use crate::compiler::ast;
4use crate::compiler::instructions::{
5    Instruction, Instructions, LocalId, LOOP_FLAG_RECURSIVE, LOOP_FLAG_WITH_LOOP_VAR, MAX_LOCALS,
6};
7use crate::compiler::tokens::Span;
8use crate::output::CaptureMode;
9use crate::value::ops::neg;
10use crate::value::{Kwargs, UndefinedType, Value, ValueMap, ValueRepr};
11
12#[cfg(test)]
13use similar_asserts::assert_eq;
14
15#[cfg(feature = "macros")]
16type Caller<'source> = ast::Spanned<ast::Macro<'source>>;
17
18#[cfg(not(feature = "macros"))]
19type Caller<'source> = std::marker::PhantomData<&'source ()>;
20
21/// For the first `MAX_LOCALS` filters/tests, an ID is returned for faster lookups from the stack.
22fn get_local_id<'source>(ids: &mut BTreeMap<&'source str, LocalId>, name: &'source str) -> LocalId {
23    if let Some(id) = ids.get(name) {
24        *id
25    } else if ids.len() >= MAX_LOCALS {
26        !0
27    } else {
28        let next_id = ids.len() as LocalId;
29        ids.insert(name, next_id);
30        next_id
31    }
32}
33
34/// Represents an open block of code that does not yet have updated
35/// jump targets.
36enum PendingBlock {
37    Branch {
38        jump_instr: u32,
39    },
40    Loop {
41        iter_instr: u32,
42        jump_instrs: Vec<u32>,
43    },
44    ScBool {
45        jump_instrs: Vec<u32>,
46    },
47}
48
49/// Provides a convenient interface to creating instructions for the VM.
50pub struct CodeGenerator<'source> {
51    instructions: Instructions<'source>,
52    blocks: BTreeMap<&'source str, Instructions<'source>>,
53    pending_block: Vec<PendingBlock>,
54    current_line: u16,
55    span_stack: Vec<Span>,
56    filter_local_ids: BTreeMap<&'source str, LocalId>,
57    test_local_ids: BTreeMap<&'source str, LocalId>,
58    raw_template_bytes: usize,
59}
60
61impl<'source> CodeGenerator<'source> {
62    /// Creates a new code generator.
63    pub fn new(file: &'source str, source: &'source str) -> CodeGenerator<'source> {
64        CodeGenerator {
65            instructions: Instructions::new(file, source),
66            blocks: BTreeMap::new(),
67            pending_block: Vec::with_capacity(32),
68            current_line: 0,
69            span_stack: Vec::with_capacity(32),
70            filter_local_ids: BTreeMap::new(),
71            test_local_ids: BTreeMap::new(),
72            raw_template_bytes: 0,
73        }
74    }
75
76    /// Sets the current location's line.
77    pub fn set_line(&mut self, lineno: u16) {
78        self.current_line = lineno;
79    }
80
81    /// Sets line from span.
82    pub fn set_line_from_span(&mut self, span: Span) {
83        self.set_line(span.start_line);
84    }
85
86    /// Pushes a span to the stack
87    pub fn push_span(&mut self, span: Span) {
88        self.span_stack.push(span);
89        self.set_line_from_span(span);
90    }
91
92    /// Pops a span from the stack.
93    pub fn pop_span(&mut self) {
94        self.span_stack.pop();
95    }
96
97    /// Add a simple instruction with the current location.
98    pub fn add(&mut self, instr: Instruction<'source>) -> u32 {
99        if let Some(span) = self.span_stack.last() {
100            if span.start_line == self.current_line {
101                return self.instructions.add_with_span(instr, *span);
102            }
103        }
104        self.instructions.add_with_line(instr, self.current_line)
105    }
106
107    /// Add a simple instruction with other location.
108    pub fn add_with_span(&mut self, instr: Instruction<'source>, span: Span) -> u32 {
109        self.instructions.add_with_span(instr, span)
110    }
111
112    /// Returns the next instruction index.
113    pub fn next_instruction(&self) -> u32 {
114        self.instructions.len() as u32
115    }
116
117    /// Creates a sub generator.
118    #[cfg(feature = "multi_template")]
119    fn new_subgenerator(&self) -> CodeGenerator<'source> {
120        let mut sub = CodeGenerator::new(self.instructions.name(), self.instructions.source());
121        sub.current_line = self.current_line;
122        sub.span_stack = self.span_stack.last().copied().into_iter().collect();
123        sub
124    }
125
126    /// Finishes a sub generator and syncs it back.
127    #[cfg(feature = "multi_template")]
128    fn finish_subgenerator(&mut self, sub: CodeGenerator<'source>) -> Instructions<'source> {
129        self.current_line = sub.current_line;
130        let (instructions, blocks) = sub.finish();
131        self.blocks.extend(blocks);
132        instructions
133    }
134
135    /// Starts a for loop
136    pub fn start_for_loop(&mut self, with_loop_var: bool, recursive: bool) {
137        let mut flags = 0;
138        if with_loop_var {
139            flags |= LOOP_FLAG_WITH_LOOP_VAR;
140        }
141        if recursive {
142            flags |= LOOP_FLAG_RECURSIVE;
143        }
144        self.add(Instruction::PushLoop(flags));
145        let instr = self.add(Instruction::Iterate(!0));
146        self.pending_block.push(PendingBlock::Loop {
147            iter_instr: instr,
148            jump_instrs: Vec::new(),
149        });
150    }
151
152    /// Ends the open for loop
153    pub fn end_for_loop(&mut self, push_did_not_iterate: bool) {
154        if let Some(PendingBlock::Loop {
155            iter_instr,
156            jump_instrs,
157        }) = self.pending_block.pop()
158        {
159            self.add(Instruction::Jump(iter_instr));
160            let loop_end = self.next_instruction();
161            if push_did_not_iterate {
162                self.add(Instruction::PushDidNotIterate);
163            };
164            self.add(Instruction::PopLoopFrame);
165            for instr in jump_instrs.into_iter().chain(Some(iter_instr)) {
166                match self.instructions.get_mut(instr) {
167                    Some(&mut Instruction::Iterate(ref mut jump_target))
168                    | Some(&mut Instruction::Jump(ref mut jump_target)) => {
169                        *jump_target = loop_end;
170                    }
171                    _ => unreachable!(),
172                }
173            }
174        } else {
175            unreachable!()
176        }
177    }
178
179    /// Begins an if conditional
180    pub fn start_if(&mut self) {
181        let jump_instr = self.add(Instruction::JumpIfFalse(!0));
182        self.pending_block.push(PendingBlock::Branch { jump_instr });
183    }
184
185    /// Begins an else conditional
186    pub fn start_else(&mut self) {
187        let jump_instr = self.add(Instruction::Jump(!0));
188        self.end_condition(jump_instr + 1);
189        self.pending_block.push(PendingBlock::Branch { jump_instr });
190    }
191
192    /// Closes the current if block.
193    pub fn end_if(&mut self) {
194        self.end_condition(self.next_instruction());
195    }
196
197    /// Starts a short-circuited bool block.
198    pub fn start_sc_bool(&mut self) {
199        self.pending_block.push(PendingBlock::ScBool {
200            jump_instrs: Vec::new(),
201        });
202    }
203
204    /// Emits a short-circuited bool operator.
205    pub fn sc_bool(&mut self, and: bool) {
206        if let Some(&mut PendingBlock::ScBool {
207            ref mut jump_instrs,
208        }) = self.pending_block.last_mut()
209        {
210            jump_instrs.push(self.instructions.add(if and {
211                Instruction::JumpIfFalseOrPop(!0)
212            } else {
213                Instruction::JumpIfTrueOrPop(!0)
214            }));
215        } else {
216            unreachable!();
217        }
218    }
219
220    /// Ends a short-circuited bool block.
221    pub fn end_sc_bool(&mut self) {
222        let end = self.next_instruction();
223        if let Some(PendingBlock::ScBool { jump_instrs }) = self.pending_block.pop() {
224            for instr in jump_instrs {
225                match self.instructions.get_mut(instr) {
226                    Some(&mut Instruction::JumpIfFalseOrPop(ref mut target))
227                    | Some(&mut Instruction::JumpIfTrueOrPop(ref mut target)) => {
228                        *target = end;
229                    }
230                    _ => unreachable!(),
231                }
232            }
233        }
234    }
235
236    fn end_condition(&mut self, new_jump_instr: u32) {
237        match self.pending_block.pop() {
238            Some(PendingBlock::Branch { jump_instr }) => {
239                match self.instructions.get_mut(jump_instr) {
240                    Some(&mut Instruction::JumpIfFalse(ref mut target))
241                    | Some(&mut Instruction::Jump(ref mut target)) => {
242                        *target = new_jump_instr;
243                    }
244                    _ => {}
245                }
246            }
247            _ => unreachable!(),
248        }
249    }
250
251    /// Compiles a statement.
252    pub fn compile_stmt(&mut self, stmt: &ast::Stmt<'source>) {
253        match stmt {
254            ast::Stmt::Template(t) => {
255                self.set_line_from_span(t.span());
256                for node in &t.children {
257                    self.compile_stmt(node);
258                }
259            }
260            ast::Stmt::EmitExpr(expr) => {
261                self.compile_emit_expr(expr);
262            }
263            ast::Stmt::EmitRaw(raw) => {
264                self.set_line_from_span(raw.span());
265                self.add(Instruction::EmitRaw(raw.raw));
266                self.raw_template_bytes += raw.raw.len();
267            }
268            ast::Stmt::ForLoop(for_loop) => {
269                self.compile_for_loop(for_loop);
270            }
271            ast::Stmt::IfCond(if_cond) => {
272                self.compile_if_stmt(if_cond);
273            }
274            ast::Stmt::WithBlock(with_block) => {
275                self.set_line_from_span(with_block.span());
276                self.add(Instruction::PushWith);
277                for (target, expr) in &with_block.assignments {
278                    self.compile_expr(expr);
279                    self.compile_assignment(target);
280                }
281                for node in &with_block.body {
282                    self.compile_stmt(node);
283                }
284                self.add(Instruction::PopFrame);
285            }
286            ast::Stmt::Set(set) => {
287                self.set_line_from_span(set.span());
288                self.compile_expr(&set.expr);
289                self.compile_assignment(&set.target);
290            }
291            ast::Stmt::SetBlock(set_block) => {
292                self.set_line_from_span(set_block.span());
293                self.add(Instruction::BeginCapture(CaptureMode::Capture));
294                for node in &set_block.body {
295                    self.compile_stmt(node);
296                }
297                self.add(Instruction::EndCapture);
298                if let Some(ref filter) = set_block.filter {
299                    self.compile_expr(filter);
300                }
301                self.compile_assignment(&set_block.target);
302            }
303            ast::Stmt::AutoEscape(auto_escape) => {
304                self.set_line_from_span(auto_escape.span());
305                self.compile_expr(&auto_escape.enabled);
306                self.add(Instruction::PushAutoEscape);
307                for node in &auto_escape.body {
308                    self.compile_stmt(node);
309                }
310                self.add(Instruction::PopAutoEscape);
311            }
312            ast::Stmt::FilterBlock(filter_block) => {
313                self.set_line_from_span(filter_block.span());
314                self.add(Instruction::BeginCapture(CaptureMode::Capture));
315                for node in &filter_block.body {
316                    self.compile_stmt(node);
317                }
318                self.add(Instruction::EndCapture);
319                self.compile_expr(&filter_block.filter);
320                self.add(Instruction::Emit);
321            }
322            #[cfg(feature = "multi_template")]
323            ast::Stmt::Block(block) => {
324                self.compile_block(block);
325            }
326            #[cfg(feature = "multi_template")]
327            ast::Stmt::Import(import) => {
328                self.add(Instruction::BeginCapture(CaptureMode::Capture));
329                self.add(Instruction::PushWith);
330                self.compile_expr(&import.expr);
331                self.add_with_span(Instruction::Include(false), import.span());
332                self.add(Instruction::EndCapture);
333                self.add(Instruction::ExportLocals);
334                self.add(Instruction::PopFrame);
335                self.compile_assignment(&import.name);
336            }
337            #[cfg(feature = "multi_template")]
338            ast::Stmt::FromImport(from_import) => {
339                self.add(Instruction::BeginCapture(CaptureMode::Discard));
340                self.add(Instruction::PushWith);
341                self.compile_expr(&from_import.expr);
342                self.add_with_span(Instruction::Include(false), from_import.span());
343                for (name, _) in &from_import.names {
344                    self.compile_expr(name);
345                }
346                self.add(Instruction::PopFrame);
347                for (name, alias) in from_import.names.iter().rev() {
348                    self.compile_assignment(alias.as_ref().unwrap_or(name));
349                }
350                self.add(Instruction::EndCapture);
351            }
352            #[cfg(feature = "multi_template")]
353            ast::Stmt::Extends(extends) => {
354                self.set_line_from_span(extends.span());
355                self.compile_expr(&extends.name);
356                self.add_with_span(Instruction::LoadBlocks, extends.span());
357            }
358            #[cfg(feature = "multi_template")]
359            ast::Stmt::Include(include) => {
360                self.set_line_from_span(include.span());
361                self.compile_expr(&include.name);
362                self.add_with_span(Instruction::Include(include.ignore_missing), include.span());
363            }
364            #[cfg(feature = "macros")]
365            ast::Stmt::Macro(macro_decl) => {
366                self.compile_macro(macro_decl);
367            }
368            #[cfg(feature = "macros")]
369            ast::Stmt::CallBlock(call_block) => {
370                self.compile_call_block(call_block);
371            }
372            #[cfg(feature = "loop_controls")]
373            ast::Stmt::Continue(cont) => {
374                self.set_line_from_span(cont.span());
375                for pending_block in self.pending_block.iter().rev() {
376                    if let PendingBlock::Loop { iter_instr, .. } = pending_block {
377                        self.add(Instruction::Jump(*iter_instr));
378                        break;
379                    }
380                }
381            }
382            #[cfg(feature = "loop_controls")]
383            ast::Stmt::Break(brk) => {
384                self.set_line_from_span(brk.span());
385                let instr = self.add(Instruction::Jump(0));
386                for pending_block in self.pending_block.iter_mut().rev() {
387                    if let &mut PendingBlock::Loop {
388                        ref mut jump_instrs,
389                        ..
390                    } = pending_block
391                    {
392                        jump_instrs.push(instr);
393                        break;
394                    }
395                }
396            }
397            ast::Stmt::Do(do_tag) => {
398                self.compile_do(do_tag);
399            }
400        }
401    }
402
403    #[cfg(feature = "multi_template")]
404    fn compile_block(&mut self, block: &ast::Spanned<ast::Block<'source>>) {
405        self.set_line_from_span(block.span());
406        let mut sub = self.new_subgenerator();
407        for node in &block.body {
408            sub.compile_stmt(node);
409        }
410        let instructions = self.finish_subgenerator(sub);
411        self.blocks.insert(block.name, instructions);
412        self.add(Instruction::CallBlock(block.name));
413    }
414
415    #[cfg(feature = "macros")]
416    fn compile_macro_expression(&mut self, macro_decl: &ast::Spanned<ast::Macro<'source>>) {
417        use crate::compiler::instructions::MACRO_CALLER;
418        self.set_line_from_span(macro_decl.span());
419        let instr = self.add(Instruction::Jump(!0));
420        let mut defaults_iter = macro_decl.defaults.iter().rev();
421        for arg in macro_decl.args.iter().rev() {
422            if let Some(default) = defaults_iter.next() {
423                self.add(Instruction::DupTop);
424                self.add(Instruction::IsUndefined);
425                self.start_if();
426                self.add(Instruction::DiscardTop);
427                self.compile_expr(default);
428                self.end_if();
429            }
430            self.compile_assignment(arg);
431        }
432        for node in &macro_decl.body {
433            self.compile_stmt(node);
434        }
435        self.add(Instruction::Return);
436        let mut undeclared = crate::compiler::meta::find_macro_closure(macro_decl);
437        let caller_reference = undeclared.remove("caller");
438        let macro_instr = self.next_instruction();
439        for name in &undeclared {
440            self.add(Instruction::Enclose(name));
441        }
442        self.add(Instruction::GetClosure);
443        self.add(Instruction::LoadConst(Value::from_object(
444            macro_decl
445                .args
446                .iter()
447                .map(|x| match x {
448                    ast::Expr::Var(var) => Value::from(var.id),
449                    _ => unreachable!(),
450                })
451                .collect::<Vec<Value>>(),
452        )));
453        let mut flags = 0;
454        if caller_reference {
455            flags |= MACRO_CALLER;
456        }
457        self.add(Instruction::BuildMacro(macro_decl.name, instr + 1, flags));
458        if let Some(&mut Instruction::Jump(ref mut target)) = self.instructions.get_mut(instr) {
459            *target = macro_instr;
460        } else {
461            unreachable!();
462        }
463    }
464
465    #[cfg(feature = "macros")]
466    fn compile_macro(&mut self, macro_decl: &ast::Spanned<ast::Macro<'source>>) {
467        self.compile_macro_expression(macro_decl);
468        self.add(Instruction::StoreLocal(macro_decl.name));
469    }
470
471    #[cfg(feature = "macros")]
472    fn compile_call_block(&mut self, call_block: &ast::Spanned<ast::CallBlock<'source>>) {
473        self.compile_call(&call_block.call, Some(&call_block.macro_decl));
474        self.add(Instruction::Emit);
475    }
476
477    fn compile_do(&mut self, do_tag: &ast::Spanned<ast::Do<'source>>) {
478        self.compile_call(&do_tag.call, None);
479    }
480
481    fn compile_if_stmt(&mut self, if_cond: &ast::Spanned<ast::IfCond<'source>>) {
482        self.set_line_from_span(if_cond.span());
483        self.push_span(if_cond.expr.span());
484        self.compile_expr(&if_cond.expr);
485        self.start_if();
486        self.pop_span();
487        for node in &if_cond.true_body {
488            self.compile_stmt(node);
489        }
490        if !if_cond.false_body.is_empty() {
491            self.start_else();
492            for node in &if_cond.false_body {
493                self.compile_stmt(node);
494            }
495        }
496        self.end_if();
497    }
498
499    fn compile_emit_expr(&mut self, expr: &ast::Spanned<ast::EmitExpr<'source>>) {
500        if let ast::Expr::Call(call) = &expr.expr {
501            self.set_line_from_span(expr.expr.span());
502            match call.identify_call() {
503                ast::CallType::Function(name) => {
504                    if name == "super" && call.args.is_empty() {
505                        self.add_with_span(Instruction::FastSuper, call.span());
506                        return;
507                    } else if name == "loop" && call.args.len() == 1 {
508                        self.compile_call_args(std::slice::from_ref(&call.args[0]), 0, None);
509                        self.add_with_span(Instruction::FastRecurse, call.span());
510                        return;
511                    }
512                }
513                #[cfg(feature = "multi_template")]
514                ast::CallType::Block(name) => {
515                    self.add(Instruction::CallBlock(name));
516                    return;
517                }
518                _ => {}
519            }
520        }
521        self.push_span(expr.expr.span());
522        self.compile_expr(&expr.expr);
523        self.add(Instruction::Emit);
524        self.pop_span();
525    }
526
527    fn compile_for_loop(&mut self, for_loop: &ast::Spanned<ast::ForLoop<'source>>) {
528        self.set_line_from_span(for_loop.span());
529
530        // filter expressions work like a nested for loop without
531        // the special loop variable. in one loop, the condition is checked and
532        // passing items accumulated into a list. in the second, that list is
533        // iterated over normally
534        if let Some(ref filter_expr) = for_loop.filter_expr {
535            self.add(Instruction::LoadConst(Value::from(0usize)));
536            self.push_span(filter_expr.span());
537            self.compile_expr(&for_loop.iter);
538            self.start_for_loop(false, false);
539            self.add(Instruction::DupTop);
540            self.compile_assignment(&for_loop.target);
541            self.compile_expr(filter_expr);
542            self.start_if();
543            self.add(Instruction::Swap);
544            self.add(Instruction::LoadConst(Value::from(1usize)));
545            self.add(Instruction::Add);
546            self.start_else();
547            self.add(Instruction::DiscardTop);
548            self.end_if();
549            self.pop_span();
550            self.end_for_loop(false);
551            self.add(Instruction::BuildList(None));
552            self.start_for_loop(true, for_loop.recursive);
553        } else {
554            self.push_span(for_loop.iter.span());
555            self.compile_expr(&for_loop.iter);
556            self.start_for_loop(true, for_loop.recursive);
557            self.pop_span();
558        }
559
560        self.compile_assignment(&for_loop.target);
561        for node in &for_loop.body {
562            self.compile_stmt(node);
563        }
564        self.end_for_loop(!for_loop.else_body.is_empty());
565        if !for_loop.else_body.is_empty() {
566            self.start_if();
567            for node in &for_loop.else_body {
568                self.compile_stmt(node);
569            }
570            self.end_if();
571        };
572    }
573
574    /// Compiles an assignment expression.
575    pub fn compile_assignment(&mut self, expr: &ast::Expr<'source>) {
576        match expr {
577            ast::Expr::Var(var) => {
578                self.add(Instruction::StoreLocal(var.id));
579            }
580            ast::Expr::List(list) => {
581                self.push_span(list.span());
582                self.add(Instruction::UnpackList(list.items.len()));
583                for expr in &list.items {
584                    self.compile_assignment(expr);
585                }
586                self.pop_span();
587            }
588            ast::Expr::GetAttr(attr) => {
589                self.push_span(attr.span());
590                self.compile_expr(&attr.expr);
591                self.add(Instruction::SetAttr(attr.name));
592            }
593            _ => unreachable!(),
594        }
595    }
596
597    /// Compiles an expression.
598    pub fn compile_expr(&mut self, expr: &ast::Expr<'source>) {
599        // try to do constant folding
600        if let Some(v) = expr.as_const() {
601            self.set_line_from_span(expr.span());
602            self.add(Instruction::LoadConst(v.clone()));
603            return;
604        }
605
606        match expr {
607            ast::Expr::Var(v) => {
608                self.set_line_from_span(v.span());
609                self.add(Instruction::Lookup(v.id));
610            }
611            ast::Expr::Const(_) => unreachable!(), // handled by constant folding
612            ast::Expr::Slice(s) => {
613                self.push_span(s.span());
614                self.compile_expr(&s.expr);
615                if let Some(ref start) = s.start {
616                    self.compile_expr(start);
617                } else {
618                    self.add(Instruction::LoadConst(Value::from(())));
619                }
620                if let Some(ref stop) = s.stop {
621                    self.compile_expr(stop);
622                } else {
623                    self.add(Instruction::LoadConst(Value::from(())));
624                }
625                if let Some(ref step) = s.step {
626                    self.compile_expr(step);
627                } else {
628                    self.add(Instruction::LoadConst(Value::from(())));
629                }
630                self.add(Instruction::Slice);
631                self.pop_span();
632            }
633            ast::Expr::UnaryOp(c) => {
634                self.set_line_from_span(c.span());
635                match c.op {
636                    ast::UnaryOpKind::Not => {
637                        self.compile_expr(&c.expr);
638                        self.add(Instruction::Not);
639                    }
640                    ast::UnaryOpKind::Neg => {
641                        // common case: negative numbers.  In that case we
642                        // directly negate them if this is possible without
643                        // an error.
644                        if let ast::Expr::Const(ref c) = c.expr {
645                            if let Ok(negated) = neg(&c.value) {
646                                self.add(Instruction::LoadConst(negated));
647                                return;
648                            }
649                        }
650                        self.compile_expr(&c.expr);
651                        self.add_with_span(Instruction::Neg, c.span());
652                    }
653                }
654            }
655            ast::Expr::BinOp(c) => {
656                self.compile_bin_op(c);
657            }
658            ast::Expr::IfExpr(i) => {
659                self.set_line_from_span(i.span());
660                self.compile_expr(&i.test_expr);
661                self.start_if();
662                self.compile_expr(&i.true_expr);
663                self.start_else();
664                if let Some(ref false_expr) = i.false_expr {
665                    self.compile_expr(false_expr);
666                } else {
667                    // special behavior: missing false block have a silent undefined
668                    // to permit special casing.  This is for compatibility also with
669                    // what Jinja2 does.
670                    self.add(Instruction::LoadConst(
671                        ValueRepr::Undefined(UndefinedType::Silent).into(),
672                    ));
673                }
674                self.end_if();
675            }
676            ast::Expr::Filter(f) => {
677                self.push_span(f.span());
678                if let Some(ref expr) = f.expr {
679                    self.compile_expr(expr);
680                }
681                let arg_count = self.compile_call_args(&f.args, 1, None);
682                let local_id = get_local_id(&mut self.filter_local_ids, f.name);
683                self.add(Instruction::ApplyFilter(f.name, arg_count, local_id));
684                self.pop_span();
685            }
686            ast::Expr::Test(f) => {
687                self.push_span(f.span());
688                self.compile_expr(&f.expr);
689                let arg_count = self.compile_call_args(&f.args, 1, None);
690                let local_id = get_local_id(&mut self.test_local_ids, f.name);
691                self.add(Instruction::PerformTest(f.name, arg_count, local_id));
692                self.pop_span();
693            }
694            ast::Expr::GetAttr(g) => {
695                self.push_span(g.span());
696                self.compile_expr(&g.expr);
697                self.add(Instruction::GetAttr(g.name));
698                self.pop_span();
699            }
700            ast::Expr::GetItem(g) => {
701                self.push_span(g.span());
702                self.compile_expr(&g.expr);
703                self.compile_expr(&g.subscript_expr);
704                self.add(Instruction::GetItem);
705                self.pop_span();
706            }
707            ast::Expr::Call(c) => {
708                self.compile_call(c, None);
709            }
710            ast::Expr::List(l) => {
711                self.set_line_from_span(l.span());
712                for item in &l.items {
713                    self.compile_expr(item);
714                }
715                self.add(Instruction::BuildList(Some(l.items.len())));
716            }
717            ast::Expr::Map(m) => {
718                self.set_line_from_span(m.span());
719                assert_eq!(m.keys.len(), m.values.len());
720                for (key, value) in m.keys.iter().zip(m.values.iter()) {
721                    self.compile_expr(key);
722                    self.compile_expr(value);
723                }
724                self.add(Instruction::BuildMap(m.keys.len()));
725            }
726        }
727    }
728
729    fn compile_call(
730        &mut self,
731        c: &ast::Spanned<ast::Call<'source>>,
732        caller: Option<&Caller<'source>>,
733    ) {
734        self.push_span(c.span());
735        match c.identify_call() {
736            ast::CallType::Function(name) => {
737                let arg_count = self.compile_call_args(&c.args, 0, caller);
738                self.add(Instruction::CallFunction(name, arg_count));
739            }
740            #[cfg(feature = "multi_template")]
741            ast::CallType::Block(name) => {
742                self.add(Instruction::BeginCapture(CaptureMode::Capture));
743                self.add(Instruction::CallBlock(name));
744                self.add(Instruction::EndCapture);
745            }
746            ast::CallType::Method(expr, name) => {
747                self.compile_expr(expr);
748                let arg_count = self.compile_call_args(&c.args, 1, caller);
749                self.add(Instruction::CallMethod(name, arg_count));
750            }
751            ast::CallType::Object(expr) => {
752                self.compile_expr(expr);
753                let arg_count = self.compile_call_args(&c.args, 1, caller);
754                self.add(Instruction::CallObject(arg_count));
755            }
756        };
757        self.pop_span();
758    }
759
760    fn compile_call_args(
761        &mut self,
762        args: &[ast::CallArg<'source>],
763        extra_args: usize,
764        caller: Option<&Caller<'source>>,
765    ) -> Option<u16> {
766        let mut pending_args = extra_args;
767        let mut num_args_batches = 0;
768        let mut has_kwargs = caller.is_some();
769        let mut static_kwargs = caller.is_none();
770
771        for arg in args {
772            match arg {
773                ast::CallArg::Pos(expr) => {
774                    self.compile_expr(expr);
775                    pending_args += 1;
776                }
777                ast::CallArg::PosSplat(expr) => {
778                    if pending_args > 0 {
779                        self.add(Instruction::BuildList(Some(pending_args)));
780                        pending_args = 0;
781                        num_args_batches += 1;
782                    }
783                    self.compile_expr(expr);
784                    num_args_batches += 1;
785                }
786                ast::CallArg::Kwarg(_, expr) => {
787                    if !matches!(expr, ast::Expr::Const(_)) {
788                        static_kwargs = false;
789                    }
790                    has_kwargs = true;
791                }
792                ast::CallArg::KwargSplat(_) => {
793                    static_kwargs = false;
794                    has_kwargs = true;
795                }
796            }
797        }
798
799        if has_kwargs {
800            let mut pending_kwargs = 0;
801            let mut num_kwargs_batches = 0;
802            let mut collected_kwargs = ValueMap::new();
803            for arg in args {
804                match arg {
805                    ast::CallArg::Kwarg(key, value) => {
806                        if static_kwargs {
807                            if let ast::Expr::Const(c) = value {
808                                collected_kwargs.insert(Value::from(*key), c.value.clone());
809                            } else {
810                                unreachable!();
811                            }
812                        } else {
813                            self.add(Instruction::LoadConst(Value::from(*key)));
814                            self.compile_expr(value);
815                            pending_kwargs += 1;
816                        }
817                    }
818                    ast::CallArg::KwargSplat(expr) => {
819                        if pending_kwargs > 0 {
820                            self.add(Instruction::BuildKwargs(pending_kwargs));
821                            num_kwargs_batches += 1;
822                            pending_kwargs = 0;
823                        }
824                        self.compile_expr(expr);
825                        num_kwargs_batches += 1;
826                    }
827                    ast::CallArg::Pos(_) | ast::CallArg::PosSplat(_) => {}
828                }
829            }
830
831            if !collected_kwargs.is_empty() {
832                self.add(Instruction::LoadConst(Kwargs::wrap(collected_kwargs)));
833            } else {
834                // The conditions above guarantee that if we collect static kwargs
835                // we cannot enter this block (single kwargs batch, no caller).
836
837                #[cfg(feature = "macros")]
838                {
839                    if let Some(caller) = caller {
840                        self.add(Instruction::LoadConst(Value::from("caller")));
841                        self.compile_macro_expression(caller);
842                        pending_kwargs += 1
843                    }
844                }
845                if num_kwargs_batches > 0 {
846                    if pending_kwargs > 0 {
847                        self.add(Instruction::BuildKwargs(pending_kwargs));
848                        num_kwargs_batches += 1;
849                    }
850                    self.add(Instruction::MergeKwargs(num_kwargs_batches));
851                } else {
852                    self.add(Instruction::BuildKwargs(pending_kwargs));
853                }
854            }
855            pending_args += 1;
856        }
857
858        if num_args_batches > 0 {
859            if pending_args > 0 {
860                self.add(Instruction::BuildList(Some(pending_args)));
861                num_args_batches += 1;
862            }
863            self.add(Instruction::UnpackLists(num_args_batches));
864            None
865        } else {
866            assert!(pending_args as u16 as usize == pending_args);
867            Some(pending_args as u16)
868        }
869    }
870
871    fn compile_bin_op(&mut self, c: &ast::Spanned<ast::BinOp<'source>>) {
872        self.push_span(c.span());
873        let instr = match c.op {
874            ast::BinOpKind::Eq => Instruction::Eq,
875            ast::BinOpKind::Ne => Instruction::Ne,
876            ast::BinOpKind::Lt => Instruction::Lt,
877            ast::BinOpKind::Lte => Instruction::Lte,
878            ast::BinOpKind::Gt => Instruction::Gt,
879            ast::BinOpKind::Gte => Instruction::Gte,
880            ast::BinOpKind::ScAnd | ast::BinOpKind::ScOr => {
881                self.start_sc_bool();
882                self.compile_expr(&c.left);
883                self.sc_bool(matches!(c.op, ast::BinOpKind::ScAnd));
884                self.compile_expr(&c.right);
885                self.end_sc_bool();
886                self.pop_span();
887                return;
888            }
889            ast::BinOpKind::Add => Instruction::Add,
890            ast::BinOpKind::Sub => Instruction::Sub,
891            ast::BinOpKind::Mul => Instruction::Mul,
892            ast::BinOpKind::Div => Instruction::Div,
893            ast::BinOpKind::FloorDiv => Instruction::IntDiv,
894            ast::BinOpKind::Rem => Instruction::Rem,
895            ast::BinOpKind::Pow => Instruction::Pow,
896            ast::BinOpKind::Concat => Instruction::StringConcat,
897            ast::BinOpKind::In => Instruction::In,
898        };
899        self.compile_expr(&c.left);
900        self.compile_expr(&c.right);
901        self.add(instr);
902        self.pop_span();
903    }
904
905    /// Returns the size hint for buffers.
906    ///
907    /// This is a proposal for the initial buffer size when rendering directly to a string.
908    pub fn buffer_size_hint(&self) -> usize {
909        // for now the assumption is made that twice the bytes of template code without
910        // control structures, rounded up to the next power of two is a good default.  The
911        // round to the next power of two is chosen because the underlying vector backing
912        // strings prefers powers of two.
913        (self.raw_template_bytes * 2).next_power_of_two()
914    }
915
916    /// Converts the compiler into the instructions.
917    pub fn finish(
918        self,
919    ) -> (
920        Instructions<'source>,
921        BTreeMap<&'source str, Instructions<'source>>,
922    ) {
923        assert!(self.pending_block.is_empty());
924        (self.instructions, self.blocks)
925    }
926}
minijinja/compiler/codegen.rs

minijinja/compiler/
codegen.rs
