Add typed block signatures and richer type mismatch displays

Author: mgomes

vibes/ast.go

@@ -277,7 +277,7 @@ func (e *RangeExpr) exprNode()     {}
 func (e *RangeExpr) Pos() Position { return e.position }
 
 type BlockLiteral struct {
-	Params   []string
+	Params   []Param
 	Body     []Statement
 	position Position
 }

vibes/execution.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"maps"
 	"math"
+	"reflect"
 	"regexp"
 	"slices"
 	"sort"
@@ -1004,7 +1005,12 @@ func (exec *Execution) CallBlock(block Value, args []Value) (Value, error) {
 		} else {
 			val = NewNil()
 		}
-		blockEnv.Define(param, val)
+		if param.Type != nil {
+			if err := checkValueType(val, param.Type); err != nil {
+				return NewNil(), exec.errorAt(param.Type.position, "%s", formatArgumentTypeMismatch(param.Name, err))
+			}
+		}
+		blockEnv.Define(param.Name, val)
 	}
 	val, returned, err := exec.evalStatements(blk.Body, blockEnv)
 	if err != nil {
@@ -3282,7 +3288,7 @@ func checkValueType(val Value, ty *TypeExpr) error {
 	}
 	return &typeMismatchError{
 		Expected: formatTypeExpr(ty),
-		Actual:   val.Kind().String(),
+		Actual:   formatValueTypeExpr(val),
 	}
 }
 
@@ -3512,6 +3518,127 @@ func formatShapeType(ty *TypeExpr) string {
 	return "{ " + strings.Join(parts, ", ") + " }"
 }
 
+func formatValueTypeExpr(val Value) string {
+	state := valueTypeFormatState{
+		seenArrays: make(map[uintptr]struct{}),
+		seenHashes: make(map[uintptr]struct{}),
+	}
+	return state.format(val)
+}
+
+type valueTypeFormatState struct {
+	seenArrays map[uintptr]struct{}
+	seenHashes map[uintptr]struct{}
+}
+
+func (s *valueTypeFormatState) format(val Value) string {
+	switch val.Kind() {
+	case KindNil:
+		return "nil"
+	case KindBool:
+		return "bool"
+	case KindInt:
+		return "int"
+	case KindFloat:
+		return "float"
+	case KindString:
+		return "string"
+	case KindMoney:
+		return "money"
+	case KindDuration:
+		return "duration"
+	case KindTime:
+		return "time"
+	case KindSymbol:
+		return "symbol"
+	case KindRange:
+		return "range"
+	case KindFunction:
+		return "function"
+	case KindBuiltin:
+		return "builtin"
+	case KindBlock:
+		return "block"
+	case KindClass:
+		return "class"
+	case KindInstance:
+		return "instance"
+	case KindArray:
+		return s.formatArray(val.Array())
+	case KindHash, KindObject:
+		return s.formatHash(val.Hash())
+	default:
+		return val.Kind().String()
+	}
+}
+
+func (s *valueTypeFormatState) formatArray(values []Value) string {
+	if len(values) == 0 {
+		return "array<empty>"
+	}
+
+	id := reflect.ValueOf(values).Pointer()
+	if id != 0 {
+		if _, seen := s.seenArrays[id]; seen {
+			return "array<...>"
+		}
+		s.seenArrays[id] = struct{}{}
+		defer delete(s.seenArrays, id)
+	}
+
+	elementTypes := make(map[string]struct{}, len(values))
+	for _, value := range values {
+		elementTypes[s.format(value)] = struct{}{}
+	}
+	return "array<" + joinSortedTypes(elementTypes) + ">"
+}
+
+func (s *valueTypeFormatState) formatHash(values map[string]Value) string {
+	if len(values) == 0 {
+		return "{}"
+	}
+
+	id := reflect.ValueOf(values).Pointer()
+	if id != 0 {
+		if _, seen := s.seenHashes[id]; seen {
+			return "{ ... }"
+		}
+		s.seenHashes[id] = struct{}{}
+		defer delete(s.seenHashes, id)
+	}
+
+	if len(values) <= 6 {
+		fields := make([]string, 0, len(values))
+		for field := range values {
+			fields = append(fields, field)
+		}
+		sort.Strings(fields)
+		parts := make([]string, len(fields))
+		for i, field := range fields {
+			parts[i] = fmt.Sprintf("%s: %s", field, s.format(values[field]))
+		}
+		return "{ " + strings.Join(parts, ", ") + " }"
+	}
+
+	valueTypes := make(map[string]struct{}, len(values))
+	for _, value := range values {
+		valueTypes[s.format(value)] = struct{}{}
+	}
+	return "hash<string, " + joinSortedTypes(valueTypes) + ">"
+}
+
+func joinSortedTypes(typeSet map[string]struct{}) string {
+	if len(typeSet) == 0 {
+		return "empty"
+	}
+	parts := make([]string, 0, len(typeSet))
+	for typeName := range typeSet {
+		parts = append(parts, typeName)
+	}
+	sort.Strings(parts)
+	return strings.Join(parts, " | ")
+}
+
 // Function looks up a compiled function by name.
 func (s *Script) Function(name string) (*ScriptFunction, bool) {
 	fn, ok := s.functions[name]

vibes/memory.go

@@ -163,7 +163,7 @@ func (est *memoryEstimator) value(val Value) int {
 		est.seenBlocks[blk] = struct{}{}
 		size += estimatedBlockBytes + estimatedSliceBaseBytes + len(blk.Params)*estimatedStringHeaderBytes
 		for _, param := range blk.Params {
-			size += len(param)
+			size += len(param.Name)
 		}
 		size += estimatedStringHeaderBytes*3 + len(blk.moduleKey) + len(blk.modulePath) + len(blk.moduleRoot)
 		size += est.env(blk.Env)

vibes/parser.go

@@ -781,7 +781,7 @@ func (p *parser) parseCallArgument(args *[]Expression, kwargs *[]KeywordArg) {
 
 func (p *parser) parseBlockLiteral() *BlockLiteral {
 	pos := p.curToken.Pos
-	params := []string{}
+	params := []Param{}
 
 	p.nextToken()
 	if p.curToken.Type == tokenPipe {
@@ -801,18 +801,18 @@ func (p *parser) parseBlockLiteral() *BlockLiteral {
 	return &BlockLiteral{Params: params, Body: body, position: pos}
 }
 
-func (p *parser) parseBlockParameters() ([]string, bool) {
-	params := []string{}
+func (p *parser) parseBlockParameters() ([]Param, bool) {
+	params := []Param{}
 	p.nextToken()
 	if p.curToken.Type == tokenPipe {
 		return params, true
 	}
 
-	if p.curToken.Type != tokenIdent {
-		p.errorExpected(p.curToken, "block parameter")
+	param, ok := p.parseBlockParameter()
+	if !ok {
 		return nil, false
 	}
-	params = append(params, p.curToken.Literal)
+	params = append(params, param)
 
 	for p.peekToken.Type == tokenComma {
 		p.nextToken()
@@ -821,11 +821,11 @@ func (p *parser) parseBlockParameters() ([]string, bool) {
 			p.addParseError(p.curToken.Pos, "trailing comma in block parameter list")
 			return nil, false
 		}
-		if p.curToken.Type != tokenIdent {
-			p.errorExpected(p.curToken, "block parameter")
+		param, ok := p.parseBlockParameter()
+		if !ok {
 			return nil, false
 		}
-		params = append(params, p.curToken.Literal)
+		params = append(params, param)
 	}
 
 	if !p.expectPeek(tokenPipe) {
@@ -835,6 +835,78 @@ func (p *parser) parseBlockParameters() ([]string, bool) {
 	return params, true
 }
 
+func (p *parser) parseBlockParameter() (Param, bool) {
+	if p.curToken.Type != tokenIdent {
+		p.errorExpected(p.curToken, "block parameter")
+		return Param{}, false
+	}
+	param := Param{Name: p.curToken.Literal}
+	if p.peekToken.Type == tokenColon {
+		p.nextToken()
+		p.nextToken()
+		param.Type = p.parseBlockParamType()
+		if param.Type == nil {
+			return Param{}, false
+		}
+	}
+	return param, true
+}
+
+func (p *parser) parseBlockParamType() *TypeExpr {
+	first := p.parseTypeAtom()
+	if first == nil {
+		return nil
+	}
+
+	union := []*TypeExpr{first}
+	for p.peekToken.Type == tokenPipe && p.blockParamUnionContinues() {
+		p.nextToken()
+		p.nextToken()
+		next := p.parseTypeAtom()
+		if next == nil {
+			return nil
+		}
+		union = append(union, next)
+	}
+
+	if len(union) == 1 {
+		return first
+	}
+
+	names := make([]string, len(union))
+	for i, option := range union {
+		names[i] = formatTypeExpr(option)
+	}
+	return &TypeExpr{
+		Name:     strings.Join(names, " | "),
+		Kind:     TypeUnion,
+		Union:    union,
+		position: first.position,
+	}
+}
+
+func (p *parser) blockParamUnionContinues() bool {
+	if p.peekToken.Type != tokenPipe {
+		return false
+	}
+
+	savedLexer := *p.l
+	savedCur := p.curToken
+	savedPeek := p.peekToken
+	savedErrors := len(p.errors)
+
+	p.nextToken()
+	p.nextToken()
+	atom := p.parseTypeAtom()
+	ok := atom != nil && (p.peekToken.Type == tokenComma || p.peekToken.Type == tokenPipe)
+
+	p.l = &savedLexer
+	p.curToken = savedCur
+	p.peekToken = savedPeek
+	p.errors = p.errors[:savedErrors]
+	return ok
+}
+
 func (p *parser) parseMemberExpression(object Expression) Expression {
 	if object == nil {
 		return nil

vibes/parser_types_test.go

@@ -81,3 +81,60 @@ end`
 		t.Fatalf("unexpected parse error: %s", got)
 	}
 }
+
+func TestParserTypeSyntaxTypedBlockParameters(t *testing.T) {
+	source := `def run(values)
+  values.map do |value: int | string, label: string?|
+    label
+  end
+end`
+
+	p := newParser(source)
+	program, errs := p.ParseProgram()
+	if len(errs) > 0 {
+		t.Fatalf("expected no parse errors, got %v", errs)
+	}
+	if len(program.Statements) != 1 {
+		t.Fatalf("expected 1 statement, got %d", len(program.Statements))
+	}
+	fn, ok := program.Statements[0].(*FunctionStmt)
+	if !ok {
+		t.Fatalf("expected function statement, got %T", program.Statements[0])
+	}
+	if len(fn.Body) != 1 {
+		t.Fatalf("expected 1 body statement, got %d", len(fn.Body))
+	}
+	exprStmt, ok := fn.Body[0].(*ExprStmt)
+	if !ok {
+		t.Fatalf("expected expression statement, got %T", fn.Body[0])
+	}
+	call, ok := exprStmt.Expr.(*CallExpr)
+	if !ok {
+		t.Fatalf("expected call expression, got %T", exprStmt.Expr)
+	}
+	if call.Block == nil {
+		t.Fatalf("expected call block")
+	}
+	if len(call.Block.Params) != 2 {
+		t.Fatalf("expected 2 block params, got %d", len(call.Block.Params))
+	}
+
+	first := call.Block.Params[0]
+	if first.Name != "value" {
+		t.Fatalf("expected first param name value, got %q", first.Name)
+	}
+	if first.Type == nil || first.Type.Kind != TypeUnion || len(first.Type.Union) != 2 {
+		t.Fatalf("expected first param union type, got %#v", first.Type)
+	}
+	if first.Type.Union[0].Kind != TypeInt || first.Type.Union[1].Kind != TypeString {
+		t.Fatalf("expected union int|string, got %#v", first.Type.Union)
+	}
+
+	second := call.Block.Params[1]
+	if second.Name != "label" {
+		t.Fatalf("expected second param name label, got %q", second.Name)
+	}
+	if second.Type == nil || second.Type.Kind != TypeString || !second.Type.Nullable {
+		t.Fatalf("expected nullable string type, got %#v", second.Type)
+	}
+}