Notebooks · lib/cas/codegen

Straight-line Viv from CAS expressions, with common subexpressions computed once

A generator states what a kernel computes as expressions in the store (lib/cas/store), one per output: S[12] = X R1 + 2 R2. This module turns the list of outputs into straight-line Viv, one assignment per line.

Common subexpressions. The store keeps one node per distinct expression, so a subexpression that two outputs share is one node with two parents. cg_emit counts the parents of every node reachable from the outputs; a node with two or more becomes a local (t1 = ...) computed once, before its first use, and every other node is written inline. Leaves (numbers, symbols) are always inline.

Products. A product in the store is flat: ex ey ez w has four factors and no node for ex ey. Two products that share a prefix of factors in the term order (ex ey ez w1 and ex ey ez w2) would each multiply it out. The generator therefore writes a product of k >= 3 factors as the product of its first k - 1 factors, made a node of its own, times the last one, so shared prefixes become shared nodes and are counted like any other. The order of the factors is the store's term order (symbols by name), which the caller uses to put the factors it wants shared first.

Numbers. An integer power b^n is written as n - 1 products; a negative one as 1.0 / (...). A rational p/q is written as the double nearest it, which is what the division p / q gives. A symbol is written as its name, unless cg_symtext gave it a text: the generators use that for numeric coefficients that are not rational (the Cartesian-to-spherical factors), which the store keeps as symbols and the kernel gets as literals.

Use.

cg_reset()
cg_out("S[3]", e1)
cg_out("S[4]", e2)
lines = cg_emit("    ")

cg_emit returns the lines, the locals first where they are needed, then the assignments in the order they were given. cg_stats() is the number of locals and of operations (each +, -, *, / counted once), for the headers the generators write.

Used with import cas.codegen.

Look inside this program

Atlas
The program’s map: every function and which calls which.
There’s no recorded run of this program yet.

Shift+Enter in a cell runs it and moves to the next. Running a cell runs every cell above it too, as one program.

[1]lines 46–373

from cas.store import (
    C_ONE, cas_add, cas_addv, cas_int, cas_isint, cas_kid, cas_kids, cas_kind, cas_mul,
    cas_mulv, cas_name_of, cas_nk, cas_np, cas_nq, cas_num, cas_pow, cas_sym, K_ADD, K_FN,
    K_MUL, K_NUM, K_POW, K_SYM,
)

# parents of each node reachable from the outputs, by node id
CG_USES = map()
# the local a shared node was written to
CG_NAME = map()
# the two factors a product of three or more is written as: prefix, last
CG_PRE = map()
CG_LAST = map()
# texts of symbols that stand for literals, by symbol name
CG_SYMTXT = map()
# the outputs (left side text, right side node) and the emitted lines, by
# position: maps, because a def cannot grow a module's array
CG_LHS = map()
CG_RHS = map()
CG_LINES = map()
# counters: outputs, locals, operations, lines
CG_COUNT = zeros(4)
CG_IN = 0
CG_LOC = 1
CG_OPS = 2
CG_NL = 3

#| Forget the outputs, the counts and the emitted locals (not the store, and
#| not the symbol texts).
def cg_reset() {
    clear(CG_USES)
    clear(CG_NAME)
    clear(CG_LHS)
    clear(CG_RHS)
    clear(CG_LINES)
    CG_COUNT[CG_IN] = 0
    CG_COUNT[CG_LOC] = 0
    CG_COUNT[CG_OPS] = 0
    CG_COUNT[CG_NL] = 0
    return 0
}

#| Write the symbol named `name` as `text` (a literal, or any Viv expression
#| that needs no parentheses in a product).
def cg_symtext(name: Str, text: Str) {
    CG_SYMTXT[name] = text
    return 0
}

#| Add an output: the line `lhs = <e>`.
def cg_out(lhs: Str, e) {
    k = CG_COUNT[CG_IN]
    CG_LHS[k] = lhs
    CG_RHS[k] = e
    CG_COUNT[CG_IN] = k + 1
    return 0
}

#| A product's non-numeric factors, in the store's order.
def cg_factors(e) {
    fs = []
    for k to cas_nk(e) {
        f = cas_kid(e, k)
        if cas_kind(f) != K_NUM {
            fs = append(fs, f)
        }
    }
    return fs
}

#| A product's numeric coefficient, `C_ONE` when it has none.
def cg_coef(e) {
    f = cas_kid(e, 0)
    if cas_kind(f) == K_NUM {
        return f
    }
    return C_ONE
}

#| The prefix and the last factor of a product of three or more factors, as
#| nodes; for a product with a coefficient, the coefficient times the rest.
def cg_split(e) {
    if has(CG_PRE, e) {
        return 0
    }
    c = cg_coef(e)
    fs = cg_factors(e)
    n = len(fs)
    if c != C_ONE {
        CG_PRE[e] = c
        if n == 1 {
            CG_LAST[e] = fs[0]
        }
        else {
            CG_LAST[e] = cas_mulv(fs)
        }
        return 0
    }
    CG_LAST[e] = fs[n - 1]
    if n == 2 {
        CG_PRE[e] = fs[0]
    }
    else {
        CG_PRE[e] = cas_mulv(fs[0:n - 1])
    }
    return 0
}

#| The children of `e` as the emitter sees them.
def cg_kids(e) {
    k = cas_kind(e)
    if k == K_MUL {
        cg_split(e)
        return [CG_PRE[e], CG_LAST[e]]
    }
    if k == K_POW {
        return [cas_kid(e, 0)]
    }
    if k == K_ADD or k == K_FN {
        return cas_kids(e)
    }
    return []
}

#| Count parents from `e` down; a node is entered once.
def cg_count(e) {
    if has(CG_USES, e) {
        CG_USES[e] = CG_USES[e] + 1
        return 0
    }
    CG_USES[e] = 1
    ks = cg_kids(e)
    for i to len(ks) {
        cg_count(ks[i])
    }
    return 0
}

#| A double as Viv source that reads back to the same double.
def cg_lit(x) {
    s = fmt(x)
    if s.find(".") < 0 and s.find("e") < 0 and s.find("inf") < 0 and s.find("nan") < 0 {
        s = s + ".0"
    }
    if x < 0.0 {
        return "(" + s + ")"
    }
    return s
}

#| The text of a leaf.
def cg_leaf(e) {
    k = cas_kind(e)
    if k == K_NUM {
        return cg_lit(cas_np(e) / cas_nq(e))
    }
    if k == K_SYM {
        nm = cas_name_of(e)
        if has(CG_SYMTXT, nm) {
            return CG_SYMTXT[nm]
        }
        return nm
    }
    stop("cg_leaf: not a leaf (a refusal reached the code generator: " + cas_name_of(e) + ")")
    return ""
}

#| The text of `e` as an operand of `*`: a sum in parentheses.
def cg_factor_text(e) {
    t = cg_text(e)
    if cas_kind(e) == K_ADD and not has(CG_NAME, e) {
        return "(" + t + ")"
    }
    return t
}

#| The text of `e`, emitting the locals it needs first. A node with two or
#| more parents is written to a local the first time and named after that.
def cg_text(e) {
    if has(CG_NAME, e) {
        return CG_NAME[e]
    }
    k = cas_kind(e)
    if k == K_NUM or k == K_SYM {
        return cg_leaf(e)
    }
    t = ""
    if k == K_MUL {
        cg_split(e)
        a = CG_PRE[e]
        b = CG_LAST[e]
        CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
        if cas_kind(a) == K_NUM and cas_np(a) == -1 and cas_nq(a) == 1 {
            t = "-" + cg_factor_text(b)
        }
        else {
            t = cg_factor_text(a) + " * " + cg_factor_text(b)
        }
    }
    if k == K_ADD {
        for i to cas_nk(e) {
            c = cas_kid(e, i)
            neg = false
            if i > 0 and cas_kind(c) == K_MUL {
                cf = cg_coef(c)
                if cas_np(cf) < 0 {
                    neg = true
                }
            }
            if neg {
                cf = cg_coef(c)
                pos = cas_mul(cas_num(-cas_np(cf), cas_nq(cf)), cas_mulv(cg_factors(c)))
                # the positive term is a new node: count it as used once here
                if not has(CG_USES, pos) {
                    CG_USES[pos] = 1
                }
                t = t + " - " + cg_factor_text(pos)
            }
            else {
                if i > 0 {
                    t = t + " + "
                }
                t = t + cg_text(c)
            }
            if i > 0 {
                CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
            }
        }
    }
    if k == K_POW {
        b = cas_kid(e, 0)
        x = cas_kid(e, 1)
        bt = cg_factor_text(b)
        if cas_isint(x) and cas_np(x) >= 2 {
            t = bt
            for i to cas_np(x) - 1 {
                t = t + " * " + bt
                CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
            }
        }
        else {
            if cas_isint(x) and cas_np(x) < 0 {
                t = bt
                for i to -cas_np(x) - 1 {
                    t = t + " * " + bt
                    CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
                }
                t = "1.0 / (" + t + ")"
                CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
            }
            else {
                t = "pow(" + bt + ", " + cg_leaf(x) + ")"
            }
        }
    }
    if k == K_FN {
        t = cas_name_of(e) + "(" + cg_text(cas_kid(e, 0)) + ")"
    }
    if CG_USES[e] >= 2 {
        CG_COUNT[CG_LOC] = CG_COUNT[CG_LOC] + 1
        nm = "t" + fmt(CG_COUNT[CG_LOC], "d")
        cg_line(CG_IND["ind"] + nm + " = " + t)
        CG_NAME[e] = nm
        return nm
    }
    return t
}

# the indentation of the lines being emitted
CG_IND = map()
CG_IND["ind"] = "    "

def cg_line(s: Str) {
    CG_LINES[CG_COUNT[CG_NL]] = s
    CG_COUNT[CG_NL] = CG_COUNT[CG_NL] + 1
    return 0
}

#| The straight-line code of the outputs given since `cg_reset`, each line
#| indented by `ind`.
def cg_emit(ind: Str) {
    CG_IND["ind"] = ind
    n = CG_COUNT[CG_IN]
    for i to n {
        cg_count(CG_RHS[i])
    }
    for i to n {
        t = cg_text(CG_RHS[i])
        cg_line(ind + CG_LHS[i] + " = " + t)
    }
    return [CG_LINES[i] for i to CG_COUNT[CG_NL]]
}

#| `[outputs, locals, operations]` of the last `cg_emit`.
def cg_stats() {
    return CG_COUNT[CG_IN:CG_OPS + 1]
}

test "a shared sum is computed once, and products share their prefixes" {
    cg_reset()
    x = cas_sym("x")
    y = cas_sym("y")
    s = cas_add(x, y)
    cg_out("a[0]", cas_mul(s, x))
    cg_out("a[1]", cas_mul(s, y))
    cg_out("a[2]", cas_mulv([x, y, cas_sym("z1")]))
    cg_out("a[3]", cas_mulv([x, y, cas_sym("z2")]))
    ls = cg_emit("")
    assert_true(ls[0] == "t1 = x + y")
    assert_true(ls[1] == "a[0] = x * t1")
    assert_true(ls[3] == "t2 = x * y")
    assert_true(ls[4] == "a[2] = t2 * z1")
    assert_true(ls[5] == "a[3] = t2 * z2")
}

test "numbers, negative terms, powers and literal symbols" {
    cg_reset()
    x = cas_sym("x")
    cg_symtext("c3", "0.5773502691896258")
    cg_out("v", cas_addv([cas_mul(cas_int(3), x), cas_mul(cas_int(-2), cas_pow(x, cas_int(2))), cas_sym("c3")]))
    cg_out("u", cas_mul(cas_num(1, 3), x))
    ls = cg_emit("")
    assert_true(ls[0] == "v = 0.5773502691896258 + 3.0 * x - 2.0 * x * x")
    assert_true(ls[1] == "u = 0.3333333333333333 * x")
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors

from cas.store import (
    C_ONE, cas_add, cas_addv, cas_int, cas_isint, cas_kid, cas_kids, cas_kind, cas_mul,
    cas_mulv, cas_name_of, cas_nk, cas_np, cas_nq, cas_num, cas_pow, cas_sym, K_ADD, K_FN,
    K_MUL, K_NUM, K_POW, K_SYM,
)

# parents of each node reachable from the outputs, by node id
CG_USES = map()
# the local a shared node was written to
CG_NAME = map()
# the two factors a product of three or more is written as: prefix, last
CG_PRE = map()
CG_LAST = map()
# texts of symbols that stand for literals, by symbol name
CG_SYMTXT = map()
# the outputs (left side text, right side node) and the emitted lines, by
# position: maps, because a def cannot grow a module's array
CG_LHS = map()
CG_RHS = map()
CG_LINES = map()
# counters: outputs, locals, operations, lines
CG_COUNT = zeros(4)
CG_IN = 0
CG_LOC = 1
CG_OPS = 2
CG_NL = 3

#| Forget the outputs, the counts and the emitted locals (not the store, and
#| not the symbol texts).
def cg_reset() {
    clear(CG_USES)
    clear(CG_NAME)
    clear(CG_LHS)
    clear(CG_RHS)
    clear(CG_LINES)
    CG_COUNT[CG_IN] = 0
    CG_COUNT[CG_LOC] = 0
    CG_COUNT[CG_OPS] = 0
    CG_COUNT[CG_NL] = 0
    return 0
}

#| Write the symbol named `name` as `text` (a literal, or any Viv expression
#| that needs no parentheses in a product).
def cg_symtext(name: Str, text: Str) {
    CG_SYMTXT[name] = text
    return 0
}

#| Add an output: the line `lhs = <e>`.
def cg_out(lhs: Str, e) {
    k = CG_COUNT[CG_IN]
    CG_LHS[k] = lhs
    CG_RHS[k] = e
    CG_COUNT[CG_IN] = k + 1
    return 0
}

#| A product's non-numeric factors, in the store's order.
def cg_factors(e) {
    fs = []
    for k to cas_nk(e) {
        f = cas_kid(e, k)
        if cas_kind(f) != K_NUM {
            fs = append(fs, f)
        }
    }
    return fs
}

#| A product's numeric coefficient, `C_ONE` when it has none.
def cg_coef(e) {
    f = cas_kid(e, 0)
    if cas_kind(f) == K_NUM {
        return f
    }
    return C_ONE
}

#| The prefix and the last factor of a product of three or more factors, as
#| nodes; for a product with a coefficient, the coefficient times the rest.
def cg_split(e) {
    if has(CG_PRE, e) {
        return 0
    }
    c = cg_coef(e)
    fs = cg_factors(e)
    n = len(fs)
    if c != C_ONE {
        CG_PRE[e] = c
        if n == 1 {
            CG_LAST[e] = fs[0]
        }
        else {
            CG_LAST[e] = cas_mulv(fs)
        }
        return 0
    }
    CG_LAST[e] = fs[n - 1]
    if n == 2 {
        CG_PRE[e] = fs[0]
    }
    else {
        CG_PRE[e] = cas_mulv(fs[0:n - 1])
    }
    return 0
}

#| The children of `e` as the emitter sees them.
def cg_kids(e) {
    k = cas_kind(e)
    if k == K_MUL {
        cg_split(e)
        return [CG_PRE[e], CG_LAST[e]]
    }
    if k == K_POW {
        return [cas_kid(e, 0)]
    }
    if k == K_ADD or k == K_FN {
        return cas_kids(e)
    }
    return []
}

#| Count parents from `e` down; a node is entered once.
def cg_count(e) {
    if has(CG_USES, e) {
        CG_USES[e] = CG_USES[e] + 1
        return 0
    }
    CG_USES[e] = 1
    ks = cg_kids(e)
    for i to len(ks) {
        cg_count(ks[i])
    }
    return 0
}

#| A double as Viv source that reads back to the same double.
def cg_lit(x) {
    s = fmt(x)
    if s.find(".") < 0 and s.find("e") < 0 and s.find("inf") < 0 and s.find("nan") < 0 {
        s = s + ".0"
    }
    if x < 0.0 {
        return "(" + s + ")"
    }
    return s
}

#| The text of a leaf.
def cg_leaf(e) {
    k = cas_kind(e)
    if k == K_NUM {
        return cg_lit(cas_np(e) / cas_nq(e))
    }
    if k == K_SYM {
        nm = cas_name_of(e)
        if has(CG_SYMTXT, nm) {
            return CG_SYMTXT[nm]
        }
        return nm
    }
    stop("cg_leaf: not a leaf (a refusal reached the code generator: " + cas_name_of(e) + ")")
    return ""
}

#| The text of `e` as an operand of `*`: a sum in parentheses.
def cg_factor_text(e) {
    t = cg_text(e)
    if cas_kind(e) == K_ADD and not has(CG_NAME, e) {
        return "(" + t + ")"
    }
    return t
}

#| The text of `e`, emitting the locals it needs first. A node with two or
#| more parents is written to a local the first time and named after that.
def cg_text(e) {
    if has(CG_NAME, e) {
        return CG_NAME[e]
    }
    k = cas_kind(e)
    if k == K_NUM or k == K_SYM {
        return cg_leaf(e)
    }
    t = ""
    if k == K_MUL {
        cg_split(e)
        a = CG_PRE[e]
        b = CG_LAST[e]
        CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
        if cas_kind(a) == K_NUM and cas_np(a) == -1 and cas_nq(a) == 1 {
            t = "-" + cg_factor_text(b)
        }
        else {
            t = cg_factor_text(a) + " * " + cg_factor_text(b)
        }
    }
    if k == K_ADD {
        for i to cas_nk(e) {
            c = cas_kid(e, i)
            neg = false
            if i > 0 and cas_kind(c) == K_MUL {
                cf = cg_coef(c)
                if cas_np(cf) < 0 {
                    neg = true
                }
            }
            if neg {
                cf = cg_coef(c)
                pos = cas_mul(cas_num(-cas_np(cf), cas_nq(cf)), cas_mulv(cg_factors(c)))
                # the positive term is a new node: count it as used once here
                if not has(CG_USES, pos) {
                    CG_USES[pos] = 1
                }
                t = t + " - " + cg_factor_text(pos)
            }
            else {
                if i > 0 {
                    t = t + " + "
                }
                t = t + cg_text(c)
            }
            if i > 0 {
                CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
            }
        }
    }
    if k == K_POW {
        b = cas_kid(e, 0)
        x = cas_kid(e, 1)
        bt = cg_factor_text(b)
        if cas_isint(x) and cas_np(x) >= 2 {
            t = bt
            for i to cas_np(x) - 1 {
                t = t + " * " + bt
                CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
            }
        }
        else {
            if cas_isint(x) and cas_np(x) < 0 {
                t = bt
                for i to -cas_np(x) - 1 {
                    t = t + " * " + bt
                    CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
                }
                t = "1.0 / (" + t + ")"
                CG_COUNT[CG_OPS] = CG_COUNT[CG_OPS] + 1
            }
            else {
                t = "pow(" + bt + ", " + cg_leaf(x) + ")"
            }
        }
    }
    if k == K_FN {
        t = cas_name_of(e) + "(" + cg_text(cas_kid(e, 0)) + ")"
    }
    if CG_USES[e] >= 2 {
        CG_COUNT[CG_LOC] = CG_COUNT[CG_LOC] + 1
        nm = "t" + fmt(CG_COUNT[CG_LOC], "d")
        cg_line(CG_IND["ind"] + nm + " = " + t)
        CG_NAME[e] = nm
        return nm
    }
    return t
}

# the indentation of the lines being emitted
CG_IND = map()
CG_IND["ind"] = "    "

def cg_line(s: Str) {
    CG_LINES[CG_COUNT[CG_NL]] = s
    CG_COUNT[CG_NL] = CG_COUNT[CG_NL] + 1
    return 0
}

#| The straight-line code of the outputs given since `cg_reset`, each line
#| indented by `ind`.
def cg_emit(ind: Str) {
    CG_IND["ind"] = ind
    n = CG_COUNT[CG_IN]
    for i to n {
        cg_count(CG_RHS[i])
    }
    for i to n {
        t = cg_text(CG_RHS[i])
        cg_line(ind + CG_LHS[i] + " = " + t)
    }
    return [CG_LINES[i] for i to CG_COUNT[CG_NL]]
}

#| `[outputs, locals, operations]` of the last `cg_emit`.
def cg_stats() {
    return CG_COUNT[CG_IN:CG_OPS + 1]
}

test "a shared sum is computed once, and products share their prefixes" {
    cg_reset()
    x = cas_sym("x")
    y = cas_sym("y")
    s = cas_add(x, y)
    cg_out("a[0]", cas_mul(s, x))
    cg_out("a[1]", cas_mul(s, y))
    cg_out("a[2]", cas_mulv([x, y, cas_sym("z1")]))
    cg_out("a[3]", cas_mulv([x, y, cas_sym("z2")]))
    ls = cg_emit("")
    assert_true(ls[0] == "t1 = x + y")
    assert_true(ls[1] == "a[0] = x * t1")
    assert_true(ls[3] == "t2 = x * y")
    assert_true(ls[4] == "a[2] = t2 * z1")
    assert_true(ls[5] == "a[3] = t2 * z2")
}

test "numbers, negative terms, powers and literal symbols" {
    cg_reset()
    x = cas_sym("x")
    cg_symtext("c3", "0.5773502691896258")
    cg_out("v", cas_addv([cas_mul(cas_int(3), x), cas_mul(cas_int(-2), cas_pow(x, cas_int(2))), cas_sym("c3")]))
    cg_out("u", cas_mul(cas_num(1, 3), x))
    ls = cg_emit("")
    assert_true(ls[0] == "v = 0.5773502691896258 + 3.0 * x - 2.0 * x * x")
    assert_true(ls[1] == "u = 0.3333333333333333 * x")
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors