Chapter 16: Environment Merging — What the Type Checker Knows After an if

In Chapter 15, we built type narrowing: inside if type(x) == "number" then, x is narrowed to Number. When the if ends, we threw away the narrowing and returned the original environment.

That’s correct for narrowing — after the if, we don’t know which branch ran. But it’s wrong for assignments. If both branches assign to the same variable, the post-if type should be the union of both branches’ types for that variable.

---@type number|string
local x = 42
if type(x) == "number" then
    x = "hello"    -- x is now String in this branch
else
    x = 99         -- x is now Number in this branch
end
-- x could be String OR Number → type is Number | String

Without merging, we’d keep the pre-if type (Number | String) or lose the assignment information entirely. Neither is right. The post-if type depends on what actually happened in each branch.

Three Cases for Merging

After an if-else, every variable falls into one of three cases:

1. Assigned in both branches → union of both types. If x becomes String in the then-block and Number in the else-block, the merged type is Number | String.

2. Assigned in only one branch → union with the pre-if type. If x becomes String in the then-block but isn’t touched in the else-block, the merged type is String | Number | String (which normalizes to Number | String). The pre-if type is still possible because the other branch might have run.

3. Not assigned in either branch → pre-if type unchanged. If nobody touched x, its type doesn’t change. Narrowing is scoped to the branch — it reverts after the if.

Case 2 is the one that surprises people. Why not use the then-block’s type? Because the else-block might have run. After an if without an else, any assignment inside the if is conditional — it might not have happened.

How Merging Works

The TypeEnv::merge function takes three environments:

• pre_env — the environment before the if-else
• then_env — the environment after the then-block
• else_env — the environment after the else-block (or None if there’s no else)

For each variable, it compares the then-type and else-type to the pre-if type. If a type changed, that branch assigned the variable. The merged type is the union of all possible types for that variable across execution paths.

#![allow(unused)]
fn main() {
fn merge(pre_env: &TypeEnv, then_env: &TypeEnv, else_env: Option<&TypeEnv>) -> TypeEnv {
    // Collect all variable names from all three environments.
    // We need every name that appears in any branch — the merged
    // environment must include variables that only one branch assigned.
    // Note: Vec::contains is O(n) per call, making this O(n²) total.
    // For the small environments in this tutorial (5-10 variables), this
    // is negligible. A production checker would use a HashMap for O(1)
    // lookups and a HashSet/BTreeSet for name deduplication.
    let mut all_names: Vec<String> = Vec::new();
    for (name, _) in &pre_env.bindings {
        if !all_names.contains(name) { all_names.push(name.clone()); }
    }
    for (name, _) in &then_env.bindings {
        if !all_names.contains(name) { all_names.push(name.clone()); }
    }
    if let Some(ee) = else_env {
        for (name, _) in &ee.bindings {
            if !all_names.contains(name) { all_names.push(name.clone()); }
        }
    }

    // For each variable, compute the merged type.
    // Key design choice: if a variable exists in a branch's environment
    // but NOT in another, the "missing" branch contributes Nil (the
    // variable might not exist if that branch ran). This is more correct
    // than degrading to Dynamic, which loses type information.
    let mut merged = TypeEnv::new();
    for name in all_names {
        let in_pre = pre_env.contains(&name);
        let in_then = then_env.contains(&name);
        let in_else = else_env.map(|e| e.contains(&name)).unwrap_or(false);

        // For each environment, use its type if the variable exists there,
        // or Nil if it doesn't (the variable is "absent" in that branch).
        let pre_type = if in_pre { pre_env.lookup(&name) } else { Type::Nil };
        let then_type = if in_then { then_env.lookup(&name) } else { Type::Nil };
        let else_type = else_env.map(|e| {
            if in_else { e.lookup(&name) } else { Type::Nil }
        });

        let then_changed = then_type != pre_type;
        let merged_type = match (then_changed, else_type) {
            // Assigned in then, no else block → union of then-type and pre-type
            (true, None) => Type::union(vec![then_type, pre_type]),
            // Assigned in then, else block exists
            (true, Some(et)) => {
                match et != pre_type {
                    // Both branches changed → union of both
                    true => Type::union(vec![then_type, et]),
                    // Only then changed → union of then-type and pre-type
                    false => Type::union(vec![then_type, pre_type]),
                }
            }
            // Not assigned in then
            (false, None) => pre_type,
            (false, Some(et)) => {
                match et != pre_type {
                    // Only else changed → union of else-type and pre-type
                    true => Type::union(vec![et, pre_type]),
                    // Neither changed → pre-if type
                    false => pre_type,
                }
            }
        };
        merged = merged.extend(name, merged_type);
    }
    merged
}
}

The key insight: we compare each branch’s final type to the pre-if type. If they differ, the variable was assigned in that branch.

The contains method

TypeEnv::lookup returns Dynamic for names that don’t exist in the environment — that’s the right behavior for type inference (“I don’t know what this is”). But for merging, Dynamic is wrong. If a variable exists in the then-block but not the else-block, using Dynamic for the else-block’s type would degrade the merged type to Dynamic (since Dynamic absorbs in unions). The right answer is Nil — the variable might not exist if that branch ran.

So merge uses a contains method to check whether a variable actually exists before looking it up:

#![allow(unused)]
fn main() {
impl TypeEnv {
    /// Check if a variable exists in this environment.
    /// Unlike `lookup` (which returns Dynamic for unknown names),
    /// this returns false for variables that don't exist.
    fn contains(&self, name: &str) -> bool {
        self.bindings.iter().rev().any(|(n, _)| n == name)
    }
}
}

When a variable is absent from a branch, merge uses Type::Nil instead of calling lookup (which would return Dynamic). This preserves type information instead of degrading it.

But wait — what about narrowing? If x starts as Number | String, narrows to Number in the then-block, and isn’t assigned, the then-block’s final type for x is Number. That’s different from the pre-if type (Number | String). Won’t the merge treat narrowing as an assignment?

Yes — and it works out correctly. The merge sees Number (then-block) vs. Number | String (pre-if), and produces Type::union(Number, Number | String) — which normalizes to Number | String. The narrowing “reverts” not because we explicitly revert it, but because the merge algorithm unions the narrowed type with the pre-if type, and union(Number, Number | String) = Number | String.

This isn’t a happy accident — it’s a natural consequence of the design. The merge’s job is to produce a type that accounts for all possible execution paths. After the if, x could be Number (if the then-block ran) or Number | String (if it didn’t). The union of those possibilities is Number | String — exactly the pre-if type. Narrowing reverts because the union operation absorbs the narrowed variant back into the original.

Narrowing + Assignment Interaction

The tricky case is when narrowing and assignment interact in the same branch:

---@type number|string
local x = 42
if type(x) == "number" then
    local y = x + 1       -- OK: x narrowed to Number
    x = "hello"           -- x is now String in this branch
else
    local z = x .. "!"    -- OK: x narrowed to String
    x = 99                -- x is now Number in this branch
end
-- After the if: x is Number | String

Inside the then-block, x is first narrowed to Number (so x + 1 works), then assigned String. The then-block’s final type for x is String. Inside the else-block, x is narrowed to String (so x .. "!" works), then assigned Number. The else-block’s final type for x is Number. The merge produces Number | String.

The order matters: narrowing happens first (applied to the branch environment before checking the body), then assignments override the narrowed type. The final type in each branch is whatever was last — narrowing or assignment.

No Else Block

When there’s no else, the merge uses None for the else environment:

---@type number|nil
local x = 42
if x ~= nil then
    local y = x + 1       -- OK: x narrowed to Number (Nil removed)
    x = "hello"           -- x is now String in this branch
end
-- x could be String (from then) or Number | Nil (pre-if, because else-path might not have run)
-- Merged type: String | Number | Nil

The None else means the else-path keeps the pre-if type. If x was assigned String in the then-block, the merged type is Type::union(String, Number | Nil) — which flattens the nested union and deduplicates, producing String | Number | Nil. The String from the then-block plus all possibilities from the pre-if type.

The Full Picture: If-Statement Type Checking

Here’s how the pieces fit together in check_stmt:

1. Evaluate the condition — infer its type and collect any diagnostics.
2. Extract narrowings from the condition expression (Chapter 15’s extract_narrowings).
3. Save the pre-if environment — we’ll need it for merging.
4. Apply narrowings for the then-block — get a narrowed environment, then check each statement in the then-block body.
5. Apply narrowings for the else-block — same process, with the complement narrowings.
6. Merge — call TypeEnv::merge(pre_env, final_then_env, final_else_env) to produce the post-if environment.

This replaces Chapter 15’s simpler approach of always returning the pre-if environment after an if. The pre-if environment is still used as the baseline for the merge — it’s the “neither branch ran” default.

What We’re Simplifying

No control-flow graph. A real type checker builds a CFG with basic blocks and computes join points; we’re doing it inline, which works for if-else but doesn’t handle loops, early returns, or break statements. In a loop body, the post-loop type should be the intersection of the pre-loop and post-loop-body types (the loop might not have run), but we return the post-loop-body type, which is overly optimistic.

No definite assignment analysis. If a variable is assigned in both branches of an if-else, a real checker can mark it as definitely assigned; we only track the type.

Narrowing reverts via union, not via scope. In a real checker, narrowing is scoped so exiting the if block restores the pre-if environment. In our implementation, narrowing “reverts” because the merge unions the narrowed type with the pre-if type, producing the original type. This works for simple cases but doesn’t generalize to nested control flow where you’d need explicit scoping.

No elseif demonstration. Lua’s elseif is syntactic sugar for else if, and our parser doesn’t produce elseif nodes (the else-block would contain another if statement). The merge handles this naturally via recursive checking, but we don’t show it.

New variables from branches. If a branch introduces a variable that doesn’t exist in the pre-if environment, the merged environment includes it with a union type including Nil — for example, if y is assigned String in the then-block but doesn’t exist in the else-block, the merged type is String | Nil (the variable might not exist if the else-block ran). We handle this by treating “variable absent from an environment” as Nil rather than Dynamic (which would degrade the type via Dynamic absorption).

Next: Chapter 17: Putting It All Together — Every concept from the tutorial in one place: annotations, unions, generics, narrowing, environment merging, and cross-file resolution. A working type checker that handles real Lua code.