dslcompile 0.0.1

High-performance symbolic mathematics with final tagless design, egglog optimization, and Rust hot-loading compilation
Documentation 
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
103
104
105
106
107

//! Pretty-printers for `ASTRepr` and `ANFExpr`
//!
//! These functions pretty-print *existing* expression trees (`ASTRepr`, `ANFExpr`),
//! using variable names from a `VariableRegistry`.
//!
//! This is different from the `PrettyPrint` final tagless instance, which builds up
//! a string as you construct an expression. These pretty-printers work on any tree,
//! regardless of how it was constructed (e.g., parsing, property-based generation, etc).

use crate::final_tagless::{ASTRepr, NumericType, VariableRegistry};
use crate::symbolic::anf::{ANFAtom, ANFComputation, ANFExpr};

/// Pretty-print an `ASTRepr` using infix notation and variable names from the registry.
pub fn pretty_ast<T: NumericType>(expr: &ASTRepr<T>, registry: &VariableRegistry) -> String {
    match expr {
        ASTRepr::Constant(val) => format!("{val}"),
        ASTRepr::Variable(idx) => {
            if *idx < registry.len() {
                registry.debug_name(*idx)
            } else {
                format!("var_{idx}")
            }
        }
        ASTRepr::Add(left, right) => {
            format!(
                "({} + {})",
                pretty_ast(left, registry),
                pretty_ast(right, registry)
            )
        }
        ASTRepr::Sub(left, right) => {
            format!(
                "({} - {})",
                pretty_ast(left, registry),
                pretty_ast(right, registry)
            )
        }
        ASTRepr::Mul(left, right) => {
            format!(
                "({} * {})",
                pretty_ast(left, registry),
                pretty_ast(right, registry)
            )
        }
        ASTRepr::Div(left, right) => {
            format!(
                "({} / {})",
                pretty_ast(left, registry),
                pretty_ast(right, registry)
            )
        }
        ASTRepr::Pow(base, exp) => {
            format!(
                "({}^{})",
                pretty_ast(base, registry),
                pretty_ast(exp, registry)
            )
        }
        ASTRepr::Neg(inner) => format!("(-{})", pretty_ast(inner, registry)),
        ASTRepr::Ln(inner) => format!("ln({})", pretty_ast(inner, registry)),
        ASTRepr::Exp(inner) => format!("exp({})", pretty_ast(inner, registry)),
        ASTRepr::Sin(inner) => format!("sin({})", pretty_ast(inner, registry)),
        ASTRepr::Cos(inner) => format!("cos({})", pretty_ast(inner, registry)),
        ASTRepr::Sqrt(inner) => format!("sqrt({})", pretty_ast(inner, registry)),
    }
}

/// Pretty-print an `ANFExpr` as indented let-bindings, using variable names from the registry.
pub fn pretty_anf<T: NumericType>(expr: &ANFExpr<T>, registry: &VariableRegistry) -> String {
    fn atom<T: NumericType>(a: &ANFAtom<T>, registry: &VariableRegistry) -> String {
        match a {
            ANFAtom::Constant(v) => format!("{v}"),
            ANFAtom::Variable(var_ref) => var_ref.debug_name(registry),
        }
    }
    fn comp<T: NumericType>(c: &ANFComputation<T>, registry: &VariableRegistry) -> String {
        match c {
            ANFComputation::Add(a, b) => format!("{} + {}", atom(a, registry), atom(b, registry)),
            ANFComputation::Sub(a, b) => format!("{} - {}", atom(a, registry), atom(b, registry)),
            ANFComputation::Mul(a, b) => format!("{} * {}", atom(a, registry), atom(b, registry)),
            ANFComputation::Div(a, b) => format!("{} / {}", atom(a, registry), atom(b, registry)),
            ANFComputation::Pow(a, b) => format!("{}^{}", atom(a, registry), atom(b, registry)),
            ANFComputation::Neg(a) => format!("-{}", atom(a, registry)),
            ANFComputation::Ln(a) => format!("ln({})", atom(a, registry)),
            ANFComputation::Exp(a) => format!("exp({})", atom(a, registry)),
            ANFComputation::Sin(a) => format!("sin({})", atom(a, registry)),
            ANFComputation::Cos(a) => format!("cos({})", atom(a, registry)),
            ANFComputation::Sqrt(a) => format!("sqrt({})", atom(a, registry)),
        }
    }
    fn go<T: NumericType>(e: &ANFExpr<T>, registry: &VariableRegistry, indent: usize) -> String {
        let spaces = "  ".repeat(indent);
        match e {
            ANFExpr::Atom(a) => atom(a, registry),
            ANFExpr::Let(var, computation, body) => {
                format!(
                    "let {} = {} in\n{}{}",
                    var.debug_name(registry),
                    comp(computation, registry),
                    spaces,
                    go(body, registry, indent)
                )
            }
        }
    }
    go(expr, registry, 0)
}