feat: Early-exit narrowing, assert narrowing and short-circuit expression narrowing

Author: HilthonTT

internal/compiler/typecheck/checker.go

@@ -283,6 +283,11 @@ func (c *checker) walkStatement(s ast.Statement) {
 		c.env.pop()
 	case *ast.ExpressionStatement:
 		c.walkExpressionDiscard(n.Expression)
+		// `assert(cond)` only returns when cond held, so its positive
+		// narrowing applies to the rest of the block.
+		if cond, ok := assertCondition(n.Expression); ok {
+			c.applyRefinement(c.refine(cond, true))
+		}
 	case *ast.DeferStatement:
 		// The deferred call is checked like any other call statement; it
 		// produces no value the surrounding scope can observe.
@@ -378,7 +383,9 @@ func (c *checker) walkAssignStatement(s *ast.AssignStatement) {
 	for i, t := range s.Targets {
 		switch tgt := t.(type) {
 		case *ast.Identifier:
-			declared, ok := c.env.lookup(tgt.Name)
+			// Check against the declared type, not any narrowing shadow —
+			// `s = nil` inside `if s ~= nil then` is legal for a `string?`.
+			declared, ok := c.env.lookupDeclared(tgt.Name)
 			if !ok {
 				// First-time global write: bind to the RHS type so later
 				// reads see it. Matches Lua's "globals materialize on
@@ -389,6 +396,9 @@ func (c *checker) walkAssignStatement(s *ast.AssignStatement) {
 			if !assignable(rhs[i], declared) {
 				c.errAssign(s.Line(), rhs[i], declared)
 			}
+			// The value changed; any active narrowing must absorb the new
+			// type or it would keep vouching for the old one.
+			c.env.widenRefined(tgt.Name, rhs[i])
 		case *ast.IndexExpression:
 			// Field assignment to a table. We currently don't enforce
 			// shape conformance on writes (most Lua tables are open),
@@ -408,7 +418,14 @@ func (c *checker) walkIfStatement(s *ast.IfStatement) {
 	// earlier condition was false — exactly Lua's evaluation order. The
 	// then-branch of each clause gets its own child frame on top.
 	c.env.push()
-	defer c.env.pop()
+
+	// Early-exit narrowing: when a leading prefix of clauses always
+	// terminates (`if s == nil then return end`), falling past the whole
+	// statement proves each of those conditions was false, so their
+	// negations outlive the `end`. Only a prefix qualifies — once a clause
+	// can fall through, later conditions may never have been evaluated.
+	var persist []refinement
+	prefixTerminates := true
 
 	for _, cl := range s.Clauses {
 		// Walk the condition once (in the already-narrowed scope) for error
@@ -423,14 +440,28 @@ func (c *checker) walkIfStatement(s *ast.IfStatement) {
 
 		// Fold this clause's "condition is false" narrowing into the
 		// accumulator so subsequent branches see it.
-		c.applyRefinement(c.refine(cl.Condition, false))
+		negR := c.refine(cl.Condition, false)
+		c.applyRefinement(negR)
+
+		if prefixTerminates {
+			if blockTerminates(cl.Body) {
+				persist = append(persist, negR)
+			} else {
+				prefixTerminates = false
+			}
+		}
 	}
 
 	if s.Else != nil {
 		c.env.push()
 		c.walkBlock(s.Else)
 		c.env.pop()
 	}
+
+	c.env.pop()
+	for _, r := range persist {
+		c.applyRefinement(r)
+	}
 }
 
 func (c *checker) walkNumericFor(s *ast.NumericForStatement) {
@@ -934,7 +965,26 @@ func (c *checker) checkNamedStructCall(call *ast.CallExpression, sc *StructCtor)
 
 func (c *checker) typeOfBinary(e *ast.BinaryExpression) *Type {
 	left := c.typeOfExpression(e.Left)
-	right := c.typeOfExpression(e.Right)
+
+	// and/or short-circuit, so their RHS only evaluates under the LHS's
+	// truthy (and) / falsy (or) outcome — type it in a frame carrying that
+	// narrowing: `s ~= nil and #s`, `x or default`.
+	var right *Type
+	switch e.Op {
+	case "and":
+		c.env.push()
+		c.applyRefinement(c.refine(e.Left, true))
+		right = c.typeOfExpression(e.Right)
+		c.env.pop()
+	case "or":
+		c.env.push()
+		c.applyRefinement(c.refine(e.Left, false))
+		right = c.typeOfExpression(e.Right)
+		c.env.pop()
+	default:
+		right = c.typeOfExpression(e.Right)
+	}
+
 	switch e.Op {
 	case "+", "-", "*", "/", "//", "%", "^", "&", "|", "~", "<<", ">>":
 		c.requireNumber(e.Line(), left, right)
@@ -965,11 +1015,22 @@ func (c *checker) typeOfBinary(e *ast.BinaryExpression) *Type {
 		}
 		return booleanT
 	case "and":
-		// `a and b` returns a if falsy, else b. The result type is the
-		// union; for a typed-correctness perspective this is good enough.
-		return NewUnion(left, right)
+		// `a and b` yields a only when a is falsy — the truthy members of
+		// a's type can't reach the result. (false can't be split off from
+		// boolean, so a boolean member survives whole.)
+		falsy := keepKinds(left, KindNil, KindBoolean)
+		if falsy.Kind == KindNever {
+			return right
+		}
+		return NewUnion(falsy, right)
 	case "or":
-		return NewUnion(left, right)
+		// `a or b` yields a only when a is truthy — nil can't survive,
+		// which is what makes `x or default` a non-optional.
+		truthy := removeKind(left, KindNil)
+		if truthy.Kind == KindNever {
+			return right
+		}
+		return NewUnion(truthy, right)
 	}
 	return anyT
 }

internal/compiler/typecheck/env.go

@@ -34,6 +34,13 @@ type env struct {
 
 type frame struct {
 	bindings map[string]*Type
+
+	// refined marks names whose binding in this frame is a narrowing shadow
+	// installed by applyRefinement, not a declaration. Assignment checking
+	// looks through these to the declared type (lookupDeclared), and an
+	// assignment widens them in place (widenRefined) so a stale narrowing
+	// can't outlive the value it described. Lazily allocated.
+	refined map[string]bool
 }
 
 func newEnv() *env {
@@ -57,12 +64,61 @@ func (e *env) pop() {
 
 // define binds a name in the innermost frame, shadowing any outer binding.
 // Re-defining within the same frame replaces the slot — matching Lua's
-// `local x ... local x` shadowing semantics.
+// `local x ... local x` shadowing semantics. A real declaration also clears
+// any refinement mark left on the slot: `local s = ...` starts fresh.
 func (e *env) define(name string, t *Type) {
 	if len(e.frames) == 0 {
 		return
 	}
-	e.frames[len(e.frames)-1].bindings[name] = t
+	f := &e.frames[len(e.frames)-1]
+	f.bindings[name] = t
+	delete(f.refined, name)
+}
+
+// defineRefined binds a narrowing shadow in the innermost frame. It types
+// exactly like define for lookup, but is invisible to lookupDeclared and
+// mutable by widenRefined.
+func (e *env) defineRefined(name string, t *Type) {
+	if len(e.frames) == 0 {
+		return
+	}
+	f := &e.frames[len(e.frames)-1]
+	f.bindings[name] = t
+	if f.refined == nil {
+		f.refined = map[string]bool{}
+	}
+	f.refined[name] = true
+}
+
+// lookupDeclared returns the innermost binding that is a declaration,
+// seeing through refinement shadows. Assignments are checked against this —
+// `s = nil` inside `if s ~= nil then` is legal because the *declared* type
+// is `string?`, whatever the branch narrowed `s` to.
+func (e *env) lookupDeclared(name string) (*Type, bool) {
+	for i := len(e.frames) - 1; i >= 0; i-- {
+		f := e.frames[i]
+		if t, ok := f.bindings[name]; ok {
+			if f.refined[name] {
+				continue
+			}
+			return t, true
+		}
+	}
+	return e.lookup(name)
+}
+
+// widenRefined folds an assigned type into every refinement shadow of
+// `name`, so a narrowing can't keep claiming the pre-assignment type after
+// `s = nil`. Widening with a union (rather than replacing) keeps outer
+// shadows sound when the assignment sits in a deeper branch that may not
+// execute on every path that reaches them.
+func (e *env) widenRefined(name string, t *Type) {
+	for i := range e.frames {
+		f := &e.frames[i]
+		if f.refined[name] {
+			f.bindings[name] = NewUnion(f.bindings[name], t)
+		}
+	}
 }
 
 // lookup walks innermost-to-outermost. Returns the bound type plus true on

internal/compiler/typecheck/refine.go

@@ -19,6 +19,18 @@ package typecheck
 //   - and / or:      conjunction propagates to the then-branch; disjunction
 //                    propagates to the else-branch (the soundly-decidable
 //                    halves of De Morgan)
+//   - early exit:    `if s == nil then return end` — when a leading prefix
+//                    of if-clauses always terminates (return/break/continue/
+//                    throw/error()), their negations persist after the `end`
+//   - assert:        `assert(cond)` applies cond's positive narrowing to the
+//                    rest of the block
+//   - expressions:   the RHS of `a and b` is typed under a's positive
+//                    narrowing; the RHS of `a or b` under its negative
+//
+// Assignment interplay: narrowing shadows are marked in the env (see
+// defineRefined) so an assignment is checked against the *declared* type and
+// then widens every shadow with the assigned type — a refinement never
+// outlives the value it described.
 //
 // Only simple identifiers are refinable — field paths like `x.y` are not
 // tracked (they can be invalidated by aliasing, and Luau itself only
@@ -195,10 +207,12 @@ func (c *checker) refineNilGuard(name string, eq bool) refinement {
 }
 
 // applyRefinement installs a refinement into the current (innermost) env
-// frame, shadowing outer bindings for the lifetime of that frame.
+// frame, shadowing outer bindings for the lifetime of that frame. The
+// bindings are marked as refinement shadows so assignment checking can see
+// through them to the declared type.
 func (c *checker) applyRefinement(r refinement) {
 	for name, t := range r {
-		c.env.define(name, t)
+		c.env.defineRefined(name, t)
 	}
 }
 
@@ -341,6 +355,80 @@ func kindForTypeString(s string) (Kind, bool) {
 	return 0, false
 }
 
+// Terminator analysis (early-exit narrowing)
+
+// blockTerminates reports whether a block never falls through: every path
+// through it ends in return/break/continue/throw or a call to error().
+// Used by walkIfStatement to persist a clause's negation past the `end` —
+// `if s == nil then return end` leaves s non-nil for the rest of the block.
+//
+// goto is deliberately NOT a terminator: its label may sit later in the
+// same block, and statements after the label would then be reached with the
+// persisted narrowing wrongly in force.
+func blockTerminates(b *ast.Block) bool {
+	if b == nil {
+		return false
+	}
+	if b.Return != nil {
+		return true
+	}
+	if len(b.Statements) == 0 {
+		return false
+	}
+	return statementTerminates(b.Statements[len(b.Statements)-1])
+}
+
+func statementTerminates(s ast.Statement) bool {
+	switch n := s.(type) {
+	case *ast.ReturnStatement, *ast.BreakStatement, *ast.ContinueStatement,
+		*ast.ThrowStatement:
+		return true
+	case *ast.ExpressionStatement:
+		return isErrorCall(n.Expression)
+	case *ast.DoStatement:
+		return blockTerminates(n.Body)
+	case *ast.IfStatement:
+		// An if terminates only when every arm does — which requires an
+		// else, or the fall-through path escapes.
+		if n.Else == nil || !blockTerminates(n.Else) {
+			return false
+		}
+		for _, cl := range n.Clauses {
+			if !blockTerminates(cl.Body) {
+				return false
+			}
+		}
+		return true
+	}
+	return false
+}
+
+// isErrorCall matches a bare `error(...)` call statement.
+func isErrorCall(e ast.Expression) bool {
+	call, ok := e.(*ast.CallExpression)
+	if !ok {
+		return false
+	}
+	fn, ok := call.Func.(*ast.Identifier)
+	return ok && fn.Name == "error"
+}
+
+// assertCondition returns the condition argument when `e` is a call of the
+// form `assert(cond, ...)`. The statement walker applies cond's positive
+// narrowing to the rest of the block: control only continues past an assert
+// that held.
+func assertCondition(e ast.Expression) (ast.Expression, bool) {
+	call, ok := e.(*ast.CallExpression)
+	if !ok || len(call.Args) < 1 {
+		return nil, false
+	}
+	fn, ok := call.Func.(*ast.Identifier)
+	if !ok || fn.Name != "assert" {
+		return nil, false
+	}
+	return call.Args[0], true
+}
+
 // AST shape helpers
 
 // typeGuardTarget returns the identifier name `x` when `e` is a call of the