redscript-mc 1.1.0 → 1.2.0

package/src/typechecker/index.ts

@@ -5,14 +5,90 @@
  * Collects errors but doesn't block compilation (warn mode).
  */
 
-import type { Program, FnDecl, Stmt, Expr, TypeNode, Block } from '../ast/types'
+import type { Program, FnDecl, Stmt, Expr, TypeNode, Block, EntityTypeName, EntitySelector } from '../ast/types'
 import { DiagnosticError, DiagnosticCollector } from '../diagnostics'
+import { getEventParamSpecs, isEventTypeName } from '../events/types'
 
 interface ScopeSymbol {
   type: TypeNode
   mutable: boolean
 }
 
+interface BuiltinSignature {
+  params: TypeNode[]
+  return: TypeNode
+}
+
+// Entity type hierarchy for subtype checking
+const ENTITY_HIERARCHY: Record<EntityTypeName, EntityTypeName | null> = {
+  'entity': null,
+  'Player': 'entity',
+  'Mob': 'entity',
+  'HostileMob': 'Mob',
+  'PassiveMob': 'Mob',
+  'Zombie': 'HostileMob',
+  'Skeleton': 'HostileMob',
+  'Creeper': 'HostileMob',
+  'Spider': 'HostileMob',
+  'Enderman': 'HostileMob',
+  'Pig': 'PassiveMob',
+  'Cow': 'PassiveMob',
+  'Sheep': 'PassiveMob',
+  'Chicken': 'PassiveMob',
+  'Villager': 'PassiveMob',
+  'ArmorStand': 'entity',
+  'Item': 'entity',
+  'Arrow': 'entity',
+}
+
+// Map Minecraft type names to entity types
+const MC_TYPE_TO_ENTITY: Record<string, EntityTypeName> = {
+  'zombie': 'Zombie',
+  'minecraft:zombie': 'Zombie',
+  'skeleton': 'Skeleton',
+  'minecraft:skeleton': 'Skeleton',
+  'creeper': 'Creeper',
+  'minecraft:creeper': 'Creeper',
+  'spider': 'Spider',
+  'minecraft:spider': 'Spider',
+  'enderman': 'Enderman',
+  'minecraft:enderman': 'Enderman',
+  'pig': 'Pig',
+  'minecraft:pig': 'Pig',
+  'cow': 'Cow',
+  'minecraft:cow': 'Cow',
+  'sheep': 'Sheep',
+  'minecraft:sheep': 'Sheep',
+  'chicken': 'Chicken',
+  'minecraft:chicken': 'Chicken',
+  'villager': 'Villager',
+  'minecraft:villager': 'Villager',
+  'armor_stand': 'ArmorStand',
+  'minecraft:armor_stand': 'ArmorStand',
+  'item': 'Item',
+  'minecraft:item': 'Item',
+  'arrow': 'Arrow',
+  'minecraft:arrow': 'Arrow',
+}
+
+const VOID_TYPE: TypeNode = { kind: 'named', name: 'void' }
+const INT_TYPE: TypeNode = { kind: 'named', name: 'int' }
+
+const BUILTIN_SIGNATURES: Record<string, BuiltinSignature> = {
+  setTimeout: {
+    params: [INT_TYPE, { kind: 'function_type', params: [], return: VOID_TYPE }],
+    return: VOID_TYPE,
+  },
+  setInterval: {
+    params: [INT_TYPE, { kind: 'function_type', params: [], return: VOID_TYPE }],
+    return: INT_TYPE,
+  },
+  clearInterval: {
+    params: [INT_TYPE],
+    return: VOID_TYPE,
+  },
+}
+
 // ---------------------------------------------------------------------------
 // Type Checker
 // ---------------------------------------------------------------------------
@@ -20,12 +96,15 @@ interface ScopeSymbol {
 export class TypeChecker {
   private collector: DiagnosticCollector
   private functions: Map<string, FnDecl> = new Map()
+  private implMethods: Map<string, Map<string, FnDecl>> = new Map()
   private structs: Map<string, Map<string, TypeNode>> = new Map()
   private enums: Map<string, Map<string, number>> = new Map()
   private consts: Map<string, TypeNode> = new Map()
   private currentFn: FnDecl | null = null
   private currentReturnType: TypeNode | null = null
   private scope: Map<string, ScopeSymbol> = new Map()
+  // Stack for tracking @s type in different contexts
+  private selfTypeStack: EntityTypeName[] = ['entity']
 
   constructor(source?: string, filePath?: string) {
     this.collector = new DiagnosticCollector(source, filePath)
@@ -53,6 +132,28 @@ export class TypeChecker {
       this.functions.set(fn.name, fn)
     }
 
+    for (const implBlock of program.implBlocks ?? []) {
+      let methods = this.implMethods.get(implBlock.typeName)
+      if (!methods) {
+        methods = new Map()
+        this.implMethods.set(implBlock.typeName, methods)
+      }
+
+      for (const method of implBlock.methods) {
+        const selfIndex = method.params.findIndex(param => param.name === 'self')
+        if (selfIndex > 0) {
+          this.report(`Method '${method.name}' must declare 'self' as the first parameter`, method.params[selfIndex])
+        }
+        if (selfIndex === 0) {
+          const selfType = this.normalizeType(method.params[0].type)
+          if (selfType.kind !== 'struct' || selfType.name !== implBlock.typeName) {
+            this.report(`Method '${method.name}' has invalid 'self' type`, method.params[0])
+          }
+        }
+        methods.set(method.name, method)
+      }
+    }
+
     for (const struct of program.structs ?? []) {
       const fields = new Map<string, TypeNode>()
       for (const field of struct.fields) {
@@ -86,6 +187,12 @@ export class TypeChecker {
       this.checkFunction(fn)
     }
 
+    for (const implBlock of program.implBlocks ?? []) {
+      for (const method of implBlock.methods) {
+        this.checkFunction(method)
+      }
+    }
+
     return this.collector.getErrors()
   }
 
@@ -95,6 +202,8 @@ export class TypeChecker {
     this.scope = new Map()
     let seenDefault = false
 
+    this.checkFunctionDecorators(fn)
+
     for (const [name, type] of this.consts.entries()) {
       this.scope.set(name, { type, mutable: false })
     }
@@ -125,6 +234,49 @@ export class TypeChecker {
     this.currentReturnType = null
   }
 
+  private checkFunctionDecorators(fn: FnDecl): void {
+    const eventDecorators = fn.decorators.filter(decorator => decorator.name === 'on')
+    if (eventDecorators.length === 0) {
+      return
+    }
+
+    if (eventDecorators.length > 1) {
+      this.report(`Function '${fn.name}' cannot have multiple @on decorators`, fn)
+      return
+    }
+
+    const eventType = eventDecorators[0].args?.eventType
+    if (!eventType) {
+      this.report(`Function '${fn.name}' is missing an event type in @on(...)`, fn)
+      return
+    }
+
+    if (!isEventTypeName(eventType)) {
+      this.report(`Unknown event type '${eventType}'`, fn)
+      return
+    }
+
+    const expectedParams = getEventParamSpecs(eventType)
+    if (fn.params.length !== expectedParams.length) {
+      this.report(
+        `Event handler '${fn.name}' for ${eventType} must declare ${expectedParams.length} parameter(s), got ${fn.params.length}`,
+        fn
+      )
+      return
+    }
+
+    for (let i = 0; i < expectedParams.length; i++) {
+      const actual = this.normalizeType(fn.params[i].type)
+      const expected = this.normalizeType(expectedParams[i].type)
+      if (!this.typesMatch(expected, actual)) {
+        this.report(
+          `Event handler '${fn.name}' parameter ${i + 1} must be ${this.typeToString(expected)}, got ${this.typeToString(actual)}`,
+          fn.params[i]
+        )
+      }
+    }
+  }
+
   private checkBlock(stmts: Block): void {
     for (const stmt of stmts) {
       this.checkStmt(stmt)
@@ -141,8 +293,7 @@ export class TypeChecker {
         break
       case 'if':
         this.checkExpr(stmt.cond)
-        this.checkBlock(stmt.then)
-        if (stmt.else_) this.checkBlock(stmt.else_)
+        this.checkIfBranches(stmt)
         break
       case 'while':
         this.checkExpr(stmt.cond)
@@ -157,7 +308,16 @@ export class TypeChecker {
       case 'foreach':
         this.checkExpr(stmt.iterable)
         if (stmt.iterable.kind === 'selector') {
-          this.scope.set(stmt.binding, { type: { kind: 'named', name: 'void' }, mutable: true }) // Entity marker
+          // Infer entity type from selector (access .sel for the EntitySelector)
+          const entityType = this.inferEntityTypeFromSelector(stmt.iterable.sel)
+          this.scope.set(stmt.binding, { 
+            type: { kind: 'entity', entityType }, 
+            mutable: false  // Entity bindings are not reassignable
+          })
+          // Push self type context for @s inside the loop
+          this.pushSelfType(entityType)
+          this.checkBlock(stmt.body)
+          this.popSelfType()
         } else {
           const iterableType = this.inferType(stmt.iterable)
           if (iterableType.kind === 'array') {
@@ -165,8 +325,8 @@ export class TypeChecker {
           } else {
             this.scope.set(stmt.binding, { type: { kind: 'named', name: 'void' }, mutable: true })
           }
+          this.checkBlock(stmt.body)
         }
-        this.checkBlock(stmt.body)
         break
       case 'match':
         this.checkExpr(stmt.expr)
@@ -180,15 +340,41 @@ export class TypeChecker {
           this.checkBlock(arm.body)
         }
         break
-      case 'as_block':
+      case 'as_block': {
+        // as block changes @s to the selector's entity type
+        const entityType = this.inferEntityTypeFromSelector(stmt.selector)
+        this.pushSelfType(entityType)
+        this.checkBlock(stmt.body)
+        this.popSelfType()
+        break
+      }
       case 'at_block':
+        // at block doesn't change @s type, only position
         this.checkBlock(stmt.body)
         break
-      case 'as_at':
+      case 'as_at': {
+        // as @x at @y - @s becomes the as selector's type
+        const entityType = this.inferEntityTypeFromSelector(stmt.as_sel)
+        this.pushSelfType(entityType)
         this.checkBlock(stmt.body)
+        this.popSelfType()
         break
+      }
       case 'execute':
+        // execute with subcommands - check for 'as' subcommands
+        for (const sub of stmt.subcommands) {
+          if (sub.kind === 'as' && sub.selector) {
+            const entityType = this.inferEntityTypeFromSelector(sub.selector)
+            this.pushSelfType(entityType)
+          }
+        }
         this.checkBlock(stmt.body)
+        // Pop for each 'as' subcommand
+        for (const sub of stmt.subcommands) {
+          if (sub.kind === 'as') {
+            this.popSelfType()
+          }
+        }
         break
       case 'expr':
         this.checkExpr(stmt.expr)
@@ -265,12 +451,24 @@ export class TypeChecker {
       case 'member':
         this.checkMemberExpr(expr)
         break
+      case 'static_call':
+        this.checkStaticCallExpr(expr)
+        break
 
       case 'binary':
         this.checkExpr(expr.left)
         this.checkExpr(expr.right)
         break
 
+      case 'is_check': {
+        this.checkExpr(expr.expr)
+        const checkedType = this.inferType(expr.expr)
+        if (checkedType.kind !== 'entity') {
+          this.report(`'is' checks require an entity expression, got ${this.typeToString(checkedType)}`, expr.expr)
+        }
+        break
+      }
+
       case 'unary':
         this.checkExpr(expr.operand)
         break
@@ -325,12 +523,6 @@ export class TypeChecker {
       case 'blockpos':
         break
 
-      case 'static_call':
-        for (const arg of expr.args) {
-          this.checkExpr(arg)
-        }
-        break
-
       // Literals don't need checking
       case 'int_lit':
       case 'float_lit':
@@ -352,6 +544,12 @@ export class TypeChecker {
       this.checkTpCall(expr)
     }
 
+    const builtin = BUILTIN_SIGNATURES[expr.fn]
+    if (builtin) {
+      this.checkFunctionCallArgs(expr.args, builtin.params, expr.fn, expr)
+      return
+    }
+
     // Check if function exists and arg count matches
     const fn = this.functions.get(expr.fn)
     if (fn) {
@@ -387,6 +585,17 @@ export class TypeChecker {
       return
     }
 
+    const implMethod = this.resolveInstanceMethod(expr)
+    if (implMethod) {
+      this.checkFunctionCallArgs(
+        expr.args,
+        implMethod.params.map(param => this.normalizeType(param.type)),
+        implMethod.name,
+        expr
+      )
+      return
+    }
+
     for (const arg of expr.args) {
       this.checkExpr(arg)
     }
@@ -497,6 +706,29 @@ export class TypeChecker {
     }
   }
 
+  private checkStaticCallExpr(expr: Extract<Expr, { kind: 'static_call' }>): void {
+    const method = this.implMethods.get(expr.type)?.get(expr.method)
+    if (!method) {
+      this.report(`Type '${expr.type}' has no static method '${expr.method}'`, expr)
+      for (const arg of expr.args) {
+        this.checkExpr(arg)
+      }
+      return
+    }
+
+    if (method.params[0]?.name === 'self') {
+      this.report(`Method '${expr.type}::${expr.method}' is an instance method`, expr)
+      return
+    }
+
+    this.checkFunctionCallArgs(
+      expr.args,
+      method.params.map(param => this.normalizeType(param.type)),
+      `${expr.type}::${expr.method}`,
+      expr
+    )
+  }
+
   private checkLambdaExpr(expr: Extract<Expr, { kind: 'lambda' }>, expectedType?: TypeNode): void {
     const normalizedExpected = expectedType ? this.normalizeType(expectedType) : undefined
     const expectedFnType = normalizedExpected?.kind === 'function_type' ? normalizedExpected : undefined
@@ -544,6 +776,42 @@ export class TypeChecker {
     this.currentReturnType = outerReturnType
   }
 
+  private checkIfBranches(stmt: Extract<Stmt, { kind: 'if' }>): void {
+    const narrowed = this.getThenBranchNarrowing(stmt.cond)
+
+    if (narrowed) {
+      const thenScope = new Map(this.scope)
+      thenScope.set(narrowed.name, { type: narrowed.type, mutable: narrowed.mutable })
+      const outerScope = this.scope
+      this.scope = thenScope
+      this.checkBlock(stmt.then)
+      this.scope = outerScope
+    } else {
+      this.checkBlock(stmt.then)
+    }
+
+    if (stmt.else_) {
+      this.checkBlock(stmt.else_)
+    }
+  }
+
+  private getThenBranchNarrowing(cond: Expr): { name: string; type: Extract<TypeNode, { kind: 'entity' }>; mutable: boolean } | null {
+    if (cond.kind !== 'is_check' || cond.expr.kind !== 'ident') {
+      return null
+    }
+
+    const symbol = this.scope.get(cond.expr.name)
+    if (!symbol || symbol.type.kind !== 'entity') {
+      return null
+    }
+
+    return {
+      name: cond.expr.name,
+      type: { kind: 'entity', entityType: cond.entityType },
+      mutable: symbol.mutable,
+    }
+  }
+
   private inferType(expr: Expr, expectedType?: TypeNode): TypeNode {
     switch (expr.kind) {
       case 'int_lit':
@@ -575,8 +843,12 @@ export class TypeChecker {
       case 'ident':
         return this.scope.get(expr.name)?.type ?? { kind: 'named', name: 'void' }
       case 'call': {
+        const builtin = BUILTIN_SIGNATURES[expr.fn]
+        if (builtin) {
+          return builtin.return
+        }
         if (expr.fn === '__array_push') {
-          return { kind: 'named', name: 'void' }
+          return VOID_TYPE
         }
         if (expr.fn === '__array_pop') {
           const target = expr.args[0]
@@ -584,20 +856,28 @@ export class TypeChecker {
             const targetType = this.scope.get(target.name)?.type
             if (targetType?.kind === 'array') return targetType.elem
           }
-          return { kind: 'named', name: 'int' }
+          return INT_TYPE
         }
         if (expr.fn === 'bossbar_get_value') {
-          return { kind: 'named', name: 'int' }
+          return INT_TYPE
         }
         if (expr.fn === 'random_sequence') {
-          return { kind: 'named', name: 'void' }
+          return VOID_TYPE
         }
         const varType = this.scope.get(expr.fn)?.type
         if (varType?.kind === 'function_type') {
           return varType.return
         }
+        const implMethod = this.resolveInstanceMethod(expr)
+        if (implMethod) {
+          return this.normalizeType(implMethod.returnType)
+        }
         const fn = this.functions.get(expr.fn)
-        return fn?.returnType ?? { kind: 'named', name: 'int' }
+        return fn?.returnType ?? INT_TYPE
+      }
+      case 'static_call': {
+        const method = this.implMethods.get(expr.type)?.get(expr.method)
+        return method ? this.normalizeType(method.returnType) : { kind: 'named', name: 'void' }
       }
       case 'invoke': {
         const calleeType = this.inferType(expr.callee)
@@ -627,6 +907,8 @@ export class TypeChecker {
           return { kind: 'named', name: 'bool' }
         }
         return this.inferType(expr.left)
+      case 'is_check':
+        return { kind: 'named', name: 'bool' }
       case 'unary':
         if (expr.op === '!') return { kind: 'named', name: 'bool' }
         return this.inferType(expr.operand)
@@ -635,6 +917,14 @@ export class TypeChecker {
           return { kind: 'array', elem: this.inferType(expr.elements[0]) }
         }
         return { kind: 'array', elem: { kind: 'named', name: 'int' } }
+      case 'struct_lit':
+        if (expectedType) {
+          const normalized = this.normalizeType(expectedType)
+          if (normalized.kind === 'struct') {
+            return normalized
+          }
+        }
+        return { kind: 'named', name: 'void' }
       case 'lambda':
         return this.inferLambdaType(
           expr,
@@ -673,6 +963,80 @@ export class TypeChecker {
     return { kind: 'function_type', params, return: returnType }
   }
 
+  // ---------------------------------------------------------------------------
+  // Entity Type Helpers
+  // ---------------------------------------------------------------------------
+
+  /** Infer entity type from a selector */
+  private inferEntityTypeFromSelector(selector: EntitySelector): EntityTypeName {
+    // @a, @p, @r always return Player
+    if (selector.kind === '@a' || selector.kind === '@p' || selector.kind === '@r') {
+      return 'Player'
+    }
+    
+    // @e or @s with type= filter
+    if (selector.filters?.type) {
+      const mcType = selector.filters.type.toLowerCase()
+      return MC_TYPE_TO_ENTITY[mcType] ?? 'entity'
+    }
+    
+    // @s uses current context
+    if (selector.kind === '@s') {
+      return this.selfTypeStack[this.selfTypeStack.length - 1]
+    }
+    
+    // Default to entity
+    return 'entity'
+  }
+
+  private resolveInstanceMethod(expr: Extract<Expr, { kind: 'call' }>): FnDecl | null {
+    const receiver = expr.args[0]
+    if (!receiver) {
+      return null
+    }
+
+    const receiverType = this.inferType(receiver)
+    if (receiverType.kind !== 'struct') {
+      return null
+    }
+
+    const method = this.implMethods.get(receiverType.name)?.get(expr.fn)
+    if (!method || method.params[0]?.name !== 'self') {
+      return null
+    }
+
+    return method
+  }
+
+  /** Check if childType is a subtype of parentType */
+  private isEntitySubtype(childType: EntityTypeName, parentType: EntityTypeName): boolean {
+    if (childType === parentType) return true
+    
+    let current: EntityTypeName | null = childType
+    while (current !== null) {
+      if (current === parentType) return true
+      current = ENTITY_HIERARCHY[current]
+    }
+    return false
+  }
+
+  /** Push a new self type context */
+  private pushSelfType(entityType: EntityTypeName): void {
+    this.selfTypeStack.push(entityType)
+  }
+
+  /** Pop self type context */
+  private popSelfType(): void {
+    if (this.selfTypeStack.length > 1) {
+      this.selfTypeStack.pop()
+    }
+  }
+
+  /** Get current @s type */
+  private getCurrentSelfType(): EntityTypeName {
+    return this.selfTypeStack[this.selfTypeStack.length - 1]
+  }
+
   private typesMatch(expected: TypeNode, actual: TypeNode): boolean {
     if (expected.kind !== actual.kind) return false
 
@@ -700,6 +1064,16 @@ export class TypeChecker {
         this.typesMatch(expected.return, actual.return)
     }
 
+    // Entity type matching with subtype support
+    if (expected.kind === 'entity' && actual.kind === 'entity') {
+      return this.isEntitySubtype(actual.entityType, expected.entityType)
+    }
+
+    // Selector matches any entity type
+    if (expected.kind === 'selector' && actual.kind === 'entity') {
+      return true
+    }
+
     return false
   }
 
@@ -715,6 +1089,12 @@ export class TypeChecker {
         return type.name
       case 'function_type':
         return `(${type.params.map(param => this.typeToString(param)).join(', ')}) -> ${this.typeToString(type.return)}`
+      case 'entity':
+        return type.entityType
+      case 'selector':
+        return 'selector'
+      default:
+        return 'unknown'
     }
   }
 
@@ -732,6 +1112,12 @@ export class TypeChecker {
     if ((type.kind === 'struct' || type.kind === 'enum') && this.enums.has(type.name)) {
       return { kind: 'enum', name: type.name }
     }
+    if (type.kind === 'struct' && type.name in ENTITY_HIERARCHY) {
+      return { kind: 'entity', entityType: type.name as EntityTypeName }
+    }
+    if (type.kind === 'named' && type.name in ENTITY_HIERARCHY) {
+      return { kind: 'entity', entityType: type.name as EntityTypeName }
+    }
     return type
   }
 }