compiler: add variable type annotations and flow-sensitive type checks

Rugo now catches more type mistakes at compile time, before you run
your program.

Variable type annotations

You can now annotate local variables, the same way you already annotate
function parameters and return types:

    name : string = "alice"
    count : int = 0

The annotation sticks for the rest of the scope. If you later try to
assign a value of the wrong type, the compiler tells you:

    count = "oops"
    # error: cannot assign string value to variable 'count' declared as int

Use `: any` when you genuinely want a variable to hold different types:

    value : any = 42
    value = "now a string"   # ok

Smarter checks across branches and reassignments

The compiler now follows the type of each variable through `if`,
`case`, `while` and `for`, so it can flag bad calls and returns even
when the offending value flows through a variable:

    def greet(name : string)
      puts "hi #{name}"
    end

    x = 42
    greet(x)
    # error: cannot pass int value as argument 1 to 'greet'
    #        (parameter 'name' declared as string)

Return statements get the same treatment:

    def lookup(id : int) : string
      if id < 0
        return id        # error: cannot return int value from
                         #        function declared returning string
      end
      "ok"
    end

When a variable could hold one of several types depending on which
branch ran, Rugo unifies them and only complains if every possible
type would be wrong. This keeps real code quiet while still catching
genuine mistakes.

Notes

- Existing programs continue to compile and run unchanged; annotations
  are opt-in.
- Error messages name the variable or parameter, the declared type,
  and the actual type, and point at the offending line.
- Documentation updated in docs/language.md.

Author: rubiojr

ast/nodes.go

@@ -283,12 +283,16 @@ type ExprStmt struct {
 func (e *ExprStmt) node() {}
 func (e *ExprStmt) stmt() {}
 
-// AssignStmt represents target = value.
+// AssignStmt represents target = value, optionally with a binding type
+// annotation (`x : int = 42`). TypeAnnot is "" when there is no
+// annotation. Annotations are sticky and validated as binding-only:
+// the parser/walker rejects them on index and field mutations.
 type AssignStmt struct {
 	BaseStmt
 	Target    string
 	Value     Expr
 	Namespace string // non-empty for top-level assignments from require'd files
+	TypeAnnot string // "" if no annotation
 }
 
 func (a *AssignStmt) node() {}

ast/var_annot_test.go

@@ -0,0 +1,85 @@
+package ast
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// findFirstAssign returns the first AssignStmt at the top level of the
+// program. Used by the variable-annotation tests below to inspect the
+// walker's output.
+func findFirstAssign(t *testing.T, prog *Program) *AssignStmt {
+	t.Helper()
+	for _, st := range prog.Statements {
+		if a, ok := st.(*AssignStmt); ok {
+			return a
+		}
+	}
+	t.Fatalf("no AssignStmt found in program")
+	return nil
+}
+
+func parseVarAnnotSource(t *testing.T, src string) (*Program, error) {
+	t.Helper()
+	c := &Compiler{}
+	return c.ParseSource(src, "test.rugo")
+}
+
+// TestVarAnnotParsesTypeAnnot exercises the `x : T = expr` syntax for
+// each of the 8 recognised type names. Each test asserts the resulting
+// AssignStmt carries the expected TypeAnnot string.
+func TestVarAnnotParsesTypeAnnot(t *testing.T) {
+	cases := []struct {
+		name string
+		src  string
+		want string
+	}{
+		{"int", "x : int = 42", "int"},
+		{"float", "x : float = 3.14", "float"},
+		{"string", `x : string = "hi"`, "string"},
+		{"bool", "x : bool = true", "bool"},
+		{"array", "x : array = [1, 2, 3]", "array"},
+		{"hash", "x : hash = {1 => 2}", "hash"},
+		{"nil", "x : nil = nil", "nil"},
+		{"any", "x : any = 42", "any"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			prog, err := parseVarAnnotSource(t, tc.src)
+			require.NoError(t, err)
+			a := findFirstAssign(t, prog)
+			assert.Equal(t, tc.want, a.TypeAnnot,
+				"AssignStmt.TypeAnnot for %q", tc.src)
+		})
+	}
+}
+
+// TestVarAnnotEmptyWithoutAnnotation asserts that `x = 42` without an
+// annotation produces AssignStmt.TypeAnnot == "".
+func TestVarAnnotEmptyWithoutAnnotation(t *testing.T) {
+	prog, err := parseVarAnnotSource(t, "x = 42")
+	require.NoError(t, err)
+	a := findFirstAssign(t, prog)
+	assert.Equal(t, "", a.TypeAnnot)
+}
+
+// TestVarAnnotRejectsIndexAndDotAssign asserts that index and field
+// assignments cannot carry a type annotation (those are mutations, not
+// bindings).
+func TestVarAnnotRejectsIndexAndDotAssign(t *testing.T) {
+	cases := []struct {
+		name string
+		src  string
+	}{
+		{"index assign", "arr = [1, 2]\narr[0] : int = 99"},
+		{"dot assign", "obj.field : int = 5"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := parseVarAnnotSource(t, tc.src)
+			assert.Error(t, err, "expected parse/walk error for %q", tc.src)
+		})
+	}
+}

ast/walker.go

@@ -1107,21 +1107,36 @@ func (w *walker) walkReturnStmt(ast []int32) (Statement, error) {
 }
 
 func (w *walker) walkAssignOrExpr(ast []int32) (Statement, error) {
-	// AssignOrExpr = Expr [ '=' Expr ] .
+	// AssignOrExpr = Expr [ ':' TypeName '=' Expr | '=' Expr ] .
 	lhs, ast, err := w.walkExpr(ast)
 	if err != nil {
 		return nil, err
 	}
 	if len(ast) > 0 {
+		// Look at the next token: either ':' (annotated binding) or '='.
+		var typeAnnot string
+		if ast[0] >= 0 {
+			tok := w.p.Token(ast[0])
+			if parser.Symbol(tok.Ch) == parser.RugoTOK_003a { // ':'
+				_, ast = w.readToken(ast) // consume ':'
+				name, rest := w.walkTypeName(ast)
+				ast = rest
+				typeAnnot = name
+			}
+		}
 		// '=' followed by Expr
 		_, ast = w.readToken(ast) // '='
 		rhs, _, err := w.walkExpr(ast)
 		if err != nil {
 			return nil, err
 		}
-		// LHS must be an identifier or index expression for assignment
+		// LHS must be an identifier or index expression for assignment.
+		// Type annotations are valid only on plain identifier bindings.
 		if ident, ok := lhs.(*IdentExpr); ok {
-			return &AssignStmt{Target: ident.Name, Value: rhs}, nil
+			return &AssignStmt{Target: ident.Name, Value: rhs, TypeAnnot: typeAnnot}, nil
+		}
+		if typeAnnot != "" {
+			return nil, &UserError{Msg: "type annotation is only valid on a plain variable binding (`x : T = expr`); index and field assignments are mutations and cannot carry an annotation"}
 		}
 		if idx, ok := lhs.(*IndexExpr); ok {
 			return &IndexAssignStmt{Object: idx.Object, Index: idx.Index, Value: rhs}, nil

compiler/check_annot.go

@@ -25,64 +25,86 @@ func TypeAnnotationCheck(sourceFile string) ast.Check {
 func (a *annotCheck) Name() string { return "type-annotation" }
 
 func (a *annotCheck) Check(prog *ast.Program) error {
+	annotated := map[string]bool{}
 	for _, s := range prog.Statements {
-		if err := a.checkStmt(s); err != nil {
+		if err := a.checkStmt(s, annotated); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func (a *annotCheck) checkStmt(s ast.Statement) error {
+func (a *annotCheck) checkStmt(s ast.Statement, annotated map[string]bool) error {
 	switch st := s.(type) {
 	case *ast.FuncDef:
 		if err := validateFuncAnnotations(st, a.sourceFile); err != nil {
 			return err
 		}
+		// Function bodies get their own annotation scope -- a `x : int = ...`
+		// inside the body is independent of any outer binding.
+		localAnnots := map[string]bool{}
 		for _, child := range st.Body {
-			if err := a.checkStmt(child); err != nil {
+			if err := a.checkStmt(child, localAnnots); err != nil {
 				return err
 			}
 		}
+	case *ast.AssignStmt:
+		if st.TypeAnnot != "" {
+			if _, ok := ParseTypeAnnotation(st.TypeAnnot); !ok {
+				return &ast.UserError{Msg: fmt.Sprintf(
+					"%s:%d: unknown type %q in annotation for variable '%s' (valid types: %s)",
+					a.sourceFile, st.SourceLine, st.TypeAnnot, st.Target, strings.Join(KnownTypeNames(), ", "),
+				)}
+			}
+			if annotated[st.Target] {
+				return &ast.UserError{Msg: fmt.Sprintf(
+					"%s:%d: re-annotation of variable '%s' (annotations are sticky bindings — assign without `: T` to update an annotated variable)",
+					a.sourceFile, st.SourceLine, st.Target,
+				)}
+			}
+			annotated[st.Target] = true
+		}
 	case *ast.TestDef:
+		localAnnots := map[string]bool{}
 		for _, child := range st.Body {
-			if err := a.checkStmt(child); err != nil {
+			if err := a.checkStmt(child, localAnnots); err != nil {
 				return err
 			}
 		}
 	case *ast.BenchDef:
+		localAnnots := map[string]bool{}
 		for _, child := range st.Body {
-			if err := a.checkStmt(child); err != nil {
+			if err := a.checkStmt(child, localAnnots); err != nil {
 				return err
 			}
 		}
 	case *ast.IfStmt:
 		for _, child := range st.Body {
-			if err := a.checkStmt(child); err != nil {
+			if err := a.checkStmt(child, annotated); err != nil {
 				return err
 			}
 		}
 		for _, c := range st.ElsifClauses {
 			for _, child := range c.Body {
-				if err := a.checkStmt(child); err != nil {
+				if err := a.checkStmt(child, annotated); err != nil {
 					return err
 				}
 			}
 		}
 		for _, child := range st.ElseBody {
-			if err := a.checkStmt(child); err != nil {
+			if err := a.checkStmt(child, annotated); err != nil {
 				return err
 			}
 		}
 	case *ast.WhileStmt:
 		for _, child := range st.Body {
-			if err := a.checkStmt(child); err != nil {
+			if err := a.checkStmt(child, annotated); err != nil {
 				return err
 			}
 		}
 	case *ast.ForStmt:
 		for _, child := range st.Body {
-			if err := a.checkStmt(child); err != nil {
+			if err := a.checkStmt(child, annotated); err != nil {
 				return err
 			}
 		}
@@ -105,8 +127,9 @@ func (a *annotCheck) checkFnExpr(fn *ast.FnExpr) error {
 	if err := validateFnExprAnnotations(fn, a.sourceFile); err != nil {
 		return err
 	}
+	localAnnots := map[string]bool{}
 	for _, child := range fn.Body {
-		if err := a.checkStmt(child); err != nil {
+		if err := a.checkStmt(child, localAnnots); err != nil {
 			return err
 		}
 	}