Repository: mitsuhiko/minijinja
Pull request #731 (2025-02-24 10:39:45)

Implement constant folding optimization
Some operations on constant values are converted into constants at codegen time.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2b575d3..9fb0083 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -32,6 +32,7 @@ All notable changes to MiniJinja are documented here.
   hidden aliases to `Function`.  #726
 - Fixed a bug that caused implicit string concatenation to not correctly
   handle escapes.  #728
+- Implemented constant folding in the code generator.  #731
 
 ## 2.8.0
 
diff --git a/minijinja/src/compiler/ast.rs b/minijinja/src/compiler/ast.rs
index 8a53493..b55e4ba 100644
--- a/minijinja/src/compiler/ast.rs
+++ b/minijinja/src/compiler/ast.rs
@@ -4,7 +4,7 @@ use std::ops::Deref;
 use std::fmt;
 
 use crate::compiler::tokens::Span;
-use crate::value::{value_map_with_capacity, Value};
+use crate::value::{ops, value_map_with_capacity, Value};
 
 /// Container for nodes with location info.
 ///
@@ -202,6 +202,47 @@ impl Expr<'_> {
             Expr::Map(s) => s.span(),
         }
     }
+
+    pub fn as_const(&self) -> Option<Value> {
+        match self {
+            Expr::Const(c) => Some(c.value.clone()),
+            Expr::List(l) => l.as_const(),
+            Expr::Map(m) => m.as_const(),
+            Expr::UnaryOp(c) => match c.op {
+                UnaryOpKind::Not => c.expr.as_const().map(|value| Value::from(!value.is_true())),
+                UnaryOpKind::Neg => c.expr.as_const().and_then(|v| ops::neg(&v).ok()),
+            },
+            Expr::BinOp(c) => {
+                let (Some(left), Some(right)) = (c.left.as_const(), c.right.as_const()) else {
+                    return None;
+                };
+                match c.op {
+                    BinOpKind::Add => ops::add(&left, &right).ok(),
+                    BinOpKind::Sub => ops::sub(&left, &right).ok(),
+                    BinOpKind::Mul => ops::mul(&left, &right).ok(),
+                    BinOpKind::Div => ops::div(&left, &right).ok(),
+                    BinOpKind::FloorDiv => ops::int_div(&left, &right).ok(),
+                    BinOpKind::Rem => ops::rem(&left, &right).ok(),
+                    BinOpKind::Pow => ops::pow(&left, &right).ok(),
+                    BinOpKind::Concat => Some(ops::string_concat(left, &right)),
+                    BinOpKind::Eq => Some(Value::from(left == right)),
+                    BinOpKind::Ne => Some(Value::from(left != right)),
+                    BinOpKind::Lt => Some(Value::from(left < right)),
+                    BinOpKind::Lte => Some(Value::from(left <= right)),
+                    BinOpKind::Gt => Some(Value::from(left > right)),
+                    BinOpKind::Gte => Some(Value::from(left >= right)),
+                    BinOpKind::In => ops::contains(&right, &left).ok(),
+                    BinOpKind::ScAnd => Some(if left.is_true() && right.is_true() {
+                        right
+                    } else {
+                        Value::from(false)
+                    }),
+                    BinOpKind::ScOr => Some(if left.is_true() { left } else { right }),
+                }
+            }
+            _ => None,
+        }
+    }
 }
 
 /// Root template node.
diff --git a/minijinja/src/compiler/codegen.rs b/minijinja/src/compiler/codegen.rs
index ba8364f..b5791d8 100644
--- a/minijinja/src/compiler/codegen.rs
+++ b/minijinja/src/compiler/codegen.rs
@@ -596,15 +596,19 @@ impl<'source> CodeGenerator<'source> {
 
     /// Compiles an expression.
     pub fn compile_expr(&mut self, expr: &ast::Expr<'source>) {
+        // try to do constant folding
+        if let Some(v) = expr.as_const() {
+            self.set_line_from_span(expr.span());
+            self.add(Instruction::LoadConst(v.clone()));
+            return;
+        }
+
         match expr {
             ast::Expr::Var(v) => {
                 self.set_line_from_span(v.span());
                 self.add(Instruction::Lookup(v.id));
             }
-            ast::Expr::Const(v) => {
-                self.set_line_from_span(v.span());
-                self.add(Instruction::LoadConst(v.value.clone()));
-            }
+            ast::Expr::Const(_) => unreachable!(), // handled by constant folding
             ast::Expr::Slice(s) => {
                 self.push_span(s.span());
                 self.compile_expr(&s.expr);


Test patch:
diff --git a/minijinja/tests/test_compiler.rs b/minijinja/tests/test_compiler.rs
index 89bdb1a..5075d84 100644
--- a/minijinja/tests/test_compiler.rs
+++ b/minijinja/tests/test_compiler.rs
@@ -1,4 +1,7 @@
 #![cfg(feature = "unstable_machinery")]
+use std::collections::BTreeMap;
+
+use minijinja::machinery::ast::Var;
 use minijinja::machinery::{CodeGenerator, Instruction};
 use minijinja::value::Value;
 
@@ -65,3 +68,139 @@ fn test_referenced_names_empty_bug() {
     let rv = instructions.get_referenced_names(0);
     assert!(rv.is_empty());
 }
+
+#[test]
+fn test_const_folding() {
+    use minijinja::machinery::{
+        ast::{BinOp, BinOpKind, Const, Expr, List, Map, Spanned, UnaryOp, UnaryOpKind},
+        Span,
+    };
+    use minijinja::Value;
+
+    // Simple constant
+    let const_expr = Expr::Const(Spanned::new(
+        Const {
+            value: Value::from(42),
+        },
+        Span::default(),
+    ));
+    assert_eq!(const_expr.as_const(), Some(Value::from(42)));
+
+    // List of constants
+    let list_expr = Expr::List(Spanned::new(
+        List {
+            items: vec![
+                Expr::Const(Spanned::new(
+                    Const {
+                        value: Value::from(1),
+                    },
+                    Span::default(),
+                )),
+                Expr::Const(Spanned::new(
+                    Const {
+                        value: Value::from(2),
+                    },
+                    Span::default(),
+                )),
+            ],
+        },
+        Span::default(),
+    ));
+    assert_eq!(
+        list_expr.as_const(),
+        Some(Value::from(vec![Value::from(1), Value::from(2)]))
+    );
+
+    // Map of constants
+    let map_expr = Expr::Map(Spanned::new(
+        Map {
+            keys: vec![Expr::Const(Spanned::new(
+                Const {
+                    value: Value::from("a"),
+                },
+                Span::default(),
+            ))],
+            values: vec![Expr::Const(Spanned::new(
+                Const {
+                    value: Value::from(1),
+                },
+                Span::default(),
+            ))],
+        },
+        Span::default(),
+    ));
+    let mut expected_map = BTreeMap::new();
+    expected_map.insert(Value::from("a"), Value::from(1));
+    assert_eq!(map_expr.as_const(), Some(Value::from(expected_map)));
+
+    // Binary op with constants
+    let binop_expr = Expr::BinOp(Spanned::new(
+        BinOp {
+            op: BinOpKind::Add,
+            left: Expr::Const(Spanned::new(
+                Const {
+                    value: Value::from(1),
+                },
+                Span::default(),
+            )),
+            right: Expr::Const(Spanned::new(
+                Const {
+                    value: Value::from(2),
+                },
+                Span::default(),
+            )),
+        },
+        Span::default(),
+    ));
+    assert_eq!(binop_expr.as_const(), Some(Value::from(3)));
+
+    // Unary op with constant
+    let unaryop_expr = Expr::UnaryOp(Spanned::new(
+        UnaryOp {
+            op: UnaryOpKind::Not,
+            expr: Expr::Const(Spanned::new(
+                Const {
+                    value: Value::from(false),
+                },
+                Span::default(),
+            )),
+        },
+        Span::default(),
+    ));
+    assert_eq!(unaryop_expr.as_const(), Some(Value::from(true)));
+
+    // Test cases that should return None
+
+    // List with var
+    let list_expr = Expr::List(Spanned::new(
+        List {
+            items: vec![
+                Expr::Const(Spanned::new(
+                    Const {
+                        value: Value::from(1),
+                    },
+                    Span::default(),
+                )),
+                Expr::Var(Spanned::new(Var { id: "foo" }, Span::default())),
+            ],
+        },
+        Span::default(),
+    ));
+    assert_eq!(list_expr.as_const(), None);
+
+    // Binary op with non-constant
+    let binop_expr = Expr::BinOp(Spanned::new(
+        BinOp {
+            op: BinOpKind::Add,
+            left: Expr::Const(Spanned::new(
+                Const {
+                    value: Value::from(1),
+                },
+                Span::default(),
+            )),
+            right: Expr::Var(Spanned::new(Var { id: "foo" }, Span::default())),
+        },
+        Span::default(),
+    ));
+    assert_eq!(binop_expr.as_const(), None);
+}