///|
/// Label information for control flow validation
priv struct LabelInfo {
  kind : LabelKind // block or loop
  block_type : @types.BlockType
}

///|
/// Kind of control structure for label resolution
priv enum LabelKind {
  BlockLabel // br jumps to end, uses results
  LoopLabel // br jumps to start, uses params
}

///|
/// Validation context for a module
priv struct ValidationContext {
  bulk_memory : Bool
  stack_switching : Bool
  types : Array[@types.SubType] // GC: now stores SubType
  funcs : Array[Int32] // type indices
  tables : Array[@types.TableType]
  mems : Array[@types.MemoryType]
  globals : Array[@types.GlobalType]
  tags : Array[@types.FuncType] // Control tags may also return values when stack switching is enabled.
  elems : Array[@types.ValueType] // element segment types
  data_count : Int32 // number of data segments
  locals : Array[@types.ValueType] // current function's locals
  labels : Array[LabelInfo] // label stack for control flow
  returns : ReadOnlyArray[@types.ValueType] // current function's return types
  local_init : Array[Bool] // tracks which locals have been initialized (for non-nullable refs)
  declared_funcs : @hashset.HashSet[Int32] // functions declared in elem/global (for ref.func validation)
  subtyping_ctx : @types.SubtypingContext // for subtype checking
}

// Type-kind lookups.
//
// Each is partial twice over: the index may be out of range, and the type it
// names may be of another kind. These used to be spelled as total functions
// returning the type directly, which left `abort` as the only way to handle
// either case — so a module naming a func type in `struct.get` killed the
// process instead of being rejected. Callers had noticed: three of them
// discarded the result and kept the call purely for that abort, which made
// the crash the validation rule.
//
// Raising says the same thing where the caller can act on it. Every caller
// already raises `ValidationError`, so nothing here needs a new channel; the
// error simply stops being fatal.

///|
/// Get the function type at `idx`, or reject the module.
fn ValidationContext::get_func_type(
  self : ValidationContext,
  idx : Int32,
) -> @types.FuncType raise ValidationError {
  if idx < 0 || idx >= self.types.length() {
    raise InvalidTypeIndex(idx)
  }
  match self.types[idx].composite {
    Func(ft) => ft
    _ => raise TypeMismatch("type at index \{idx} is not a function type")
  }
}

///|
/// Get the struct type at `idx`, or reject the module.
fn ValidationContext::get_struct_type(
  self : ValidationContext,
  idx : Int32,
) -> @types.StructType raise ValidationError {
  if idx < 0 || idx >= self.types.length() {
    raise InvalidTypeIndex(idx)
  }
  match self.types[idx].composite {
    Struct(st) => st
    _ => raise TypeMismatch("type at index \{idx} is not a struct type")
  }
}

///|
/// Get the array type at `idx`, or reject the module.
fn ValidationContext::get_array_type(
  self : ValidationContext,
  idx : Int32,
) -> @types.ArrayType raise ValidationError {
  if idx < 0 || idx >= self.types.length() {
    raise InvalidTypeIndex(idx)
  }
  match self.types[idx].composite {
    Array(at) => at
    _ => raise TypeMismatch("type at index \{idx} is not an array type")
  }
}

///|
/// Convert StorageType to ValueType for validation
fn storage_type_to_value_type(
  storage_type : @types.StorageType,
) -> @types.ValueType {
  match storage_type {
    Val(vt) => vt
    Packed(_) => I32 // Packed types are accessed as i32
  }
}

///|
/// Collect all value types referenced in a SubType (for validation)
fn collect_value_types_from_subtype(
  subtype : @types.SubType,
) -> Array[@types.ValueType] {
  let result : Array[@types.ValueType] = []
  match subtype.composite {
    Cont(index) => result.push(RefFuncTyped(index))
    Func(ft) => {
      for p in ft.params {
        result.push(p)
      }
      for r in ft.results {
        result.push(r)
      }
    }
    Struct(st) =>
      for field in st.fields {
        match field.storage_type {
          Val(vt) => result.push(vt)
          Packed(_) => () // i8/i16 don't contain type indices
        }
      }
    Array(at) =>
      match at.element.storage_type {
        Val(vt) => result.push(vt)
        Packed(_) => ()
      }
  }
  result
}

///|
/// Collect function indices referenced by ref.func in an instruction sequence
fn collect_declared_funcs(
  instrs : ArrayView[@types.Instruction],
  declared : @hashset.HashSet[Int32],
) -> Unit {
  for instr in instrs {
    if instr is RefFunc(func_idx) {
      declared.add(func_idx)
    }
  }
}

///|
fn ValidationContext::ValidationContext(
  mod : @types.Module,
  bulk_memory~ : Bool = true,
  stack_switching~ : Bool = false,
) -> ValidationContext raise ValidationError {
  if mod.funcs.length() != mod.codes.length() {
    raise TypeMismatch("function and code section lengths differ")
  }
  if !mod.type_rec_groups.is_empty() &&
    mod.type_rec_groups.length() != mod.types.length() {
    raise TypeMismatch("recursive group map length differs from type section")
  }
  // A recursive group occupies one contiguous range in the type index space.
  let groups : @hashset.HashSet[Int32] = HashSet([])
  let mut previous_group = -1
  for group in mod.type_rec_groups {
    guard group >= 0 else { raise TypeMismatch("negative recursive group ID") }
    if group != previous_group {
      guard !groups.contains(group) else {
        raise TypeMismatch("non-contiguous recursive group")
      }
      groups.add(group)
      previous_group = group
    }
  }
  for index, subtype in mod.types {
    guard subtype.supertypes.length() <= 1 else {
      raise TypeMismatch("multiple supertypes are not supported by Core Wasm")
    }
    for parent in subtype.supertypes {
      guard parent >= 0 && parent < index else {
        raise TypeMismatch("supertypes must be defined before subtypes")
      }
    }
    for type_ in collect_value_types_from_subtype(subtype) {
      validate_value_type(type_, mod.types, stack_switching~)
      validate_type_forward_refs(type_, index, mod.type_rec_groups)
    }
  }
  for data in mod.datas {
    guard data.memory_idx >= -1 else {
      raise InvalidMemoryIndex(data.memory_idx)
    }
  }
  // Collect function types (imports + defined functions)
  let funcs : Array[Int32] = []
  for imp in mod.imports {
    if imp.desc is Func(type_idx) {
      funcs.push(type_idx)
    }
  }
  for type_idx in mod.funcs {
    funcs.push(type_idx)
  }
  for index in funcs {
    guard index >= 0 && index < mod.types.length() else {
      raise InvalidTypeIndex(index)
    }
    guard mod.types[index].composite is Func(_) else {
      raise TypeMismatch("function declaration must name a function type")
    }
  }
  // Collect tables
  let tables : Array[@types.TableType] = []
  for imp in mod.imports {
    if imp.desc is Table(table_type) {
      tables.push(table_type)
    }
  }
  for table in mod.tables {
    tables.push(table.type_)
  }
  // Collect memories
  let mems : Array[@types.MemoryType] = []
  for imp in mod.imports {
    if imp.desc is Memory(mem_type) {
      mems.push(mem_type)
    }
  }
  for mem in mod.memories {
    mems.push(mem)
  }
  // Collect globals
  let globals : Array[@types.GlobalType] = []
  for imp in mod.imports {
    if imp.desc is Global(global_type) {
      globals.push(global_type)
    }
  }
  for global in mod.globals {
    globals.push(global.type_)
  }
  for global in globals {
    validate_value_type(global.value_type, mod.types, stack_switching~)
  }
  for table in tables {
    validate_value_type(table.elem_type, mod.types, stack_switching~)
    guard (match table.elem_type {
      I32 | I64 | F32 | F64 | V128 => false
      _ => true
    }) else {
      raise TypeMismatch("table element type must be a reference")
    }
    validate_table_limits(table.limits, is_table64=table.is_table64)
  }
  for element in mod.elems {
    validate_value_type(element.type_, mod.types, stack_switching~)
  }
  // Collect elem segment types
  let elems : Array[@types.ValueType] = []
  for elem in mod.elems {
    elems.push(elem.type_)
  }
  // Collect tags (imports + defined tags)
  // A tag names a function type. One that names a struct or array is a
  // module to reject here, where the context is still being built, rather
  // than a crash on the way in.
  let tags : Array[@types.FuncType] = []
  fn tag_type(type_idx : Int32) -> @types.FuncType raise ValidationError {
    guard mod.func_type_at(type_idx) is Some(ft) else {
      raise TypeMismatch("tag type at index \{type_idx} is not a function type")
    }
    if !stack_switching && !ft.results.is_empty() {
      raise TypeMismatch("tag results require stack switching")
    }
    ft
  }

  for imp in mod.imports {
    if imp.desc is Tag(type_idx) {
      tags.push(tag_type(type_idx))
    }
  }
  for tag in mod.tags {
    tags.push(tag_type(tag.type_idx))
  }

  // Collect declared functions from elem segments and global inits
  let declared_funcs : @hashset.HashSet[Int32] = HashSet([])
  // Functions referenced in element segments
  for elem in mod.elems {
    for init in elem.init {
      collect_declared_funcs(init, declared_funcs)
    }
  }
  // Functions referenced in global initializers
  for global in mod.globals {
    collect_declared_funcs(global.init, declared_funcs)
  }
  // Exported functions are declared and can be referenced by ref.func.
  for exp in mod.exports {
    if exp.desc is Func(func_idx) {
      declared_funcs.add(func_idx)
    }
  }
  let canonical = @types.compute_canonical_type_indices(
    mod.types,
    type_rec_groups=mod.type_rec_groups,
  )
  {
    types: mod.types,
    funcs,
    tables,
    mems,
    globals,
    tags,
    elems,
    bulk_memory,
    stack_switching,
    data_count: mod.datas.length(),
    locals: [],
    labels: [],
    returns: [],
    local_init: [],
    declared_funcs,
    subtyping_ctx: @types.SubtypingContext::same_module(
      mod.types,
      rec_groups=mod.type_rec_groups,
      canonical~,
    ),
  }
}

///|
/// Check if a value type is a non-nullable reference type
fn is_non_nullable_ref(ty : @types.ValueType) -> Bool {
  match ty {
    RefFunc
    | RefExtern
    | RefFuncTyped(_)
    | RefCont
    | RefNoCont
    | RefContTyped(_) => true
    _ => false
  }
}

///|
/// Check if a value type is any reference type (nullable or non-nullable)
fn is_ref_type(ty : @types.ValueType) -> Bool {
  match ty {
    FuncRef
    | ExternRef
    | RefFunc
    | RefExtern
    | RefFuncTyped(_)
    | RefNullFuncTyped(_)
    | ContRef
    | RefCont
    | RefNoCont
    | NullContRef
    | RefContTyped(_)
    | RefNullContTyped(_) => true
    _ => false
  }
}

///|
/// Validate that any type indices in a value type are valid
fn validate_value_type(
  ty : @types.ValueType,
  types : ArrayView[@types.SubType],
  stack_switching~ : Bool = false,
) -> Unit raise ValidationError {
  if ty.is_continuation_ref() && !stack_switching {
    raise TypeMismatch("stack switching is not enabled")
  }
  let idx : Int32? = match ty {
    RefFuncTyped(i)
    | RefNullFuncTyped(i)
    | RefContTyped(i)
    | RefNullContTyped(i) => Some(i)
    // Ref(Null)Struct/Array(-1) encode abstract heap types (struct/array),
    // which do not refer to a concrete type index.
    RefStruct(i) | RefNullStruct(i) => if i == -1 { None } else { Some(i) }
    RefArray(i) | RefNullArray(i) => if i == -1 { None } else { Some(i) }
    _ => None
  }
  if idx is Some(i) && (i < 0 || i >= types.length()) {
    raise InvalidTypeIndex(i)
  }
  if idx is Some(i) {
    let matches = match (ty, types[i].composite) {
      (RefFuncTyped(_) | RefNullFuncTyped(_), Func(_)) => true
      (RefContTyped(_) | RefNullContTyped(_), Cont(_)) => true
      (RefStruct(_) | RefNullStruct(_), Struct(_)) => true
      (RefArray(_) | RefNullArray(_), Array(_)) => true
      _ => false
    }
    if !matches {
      raise TypeMismatch("reference type does not match the indexed type kind")
    }
  }
}

///|
/// Validate that forward type references are only allowed within rec groups.
/// For a type at index `type_idx`, any reference to a type with higher index
/// must be in the same rec group.
fn validate_type_forward_refs(
  ty : @types.ValueType,
  type_idx : Int32,
  type_rec_groups : Array[Int32],
) -> Unit raise ValidationError {
  let ref_idx : Int32? = match ty {
    RefFuncTyped(i)
    | RefNullFuncTyped(i)
    | RefContTyped(i)
    | RefNullContTyped(i) => Some(i)
    // Ref(Null)Struct/Array(-1) are abstract heap types, not forward refs.
    RefStruct(i) | RefNullStruct(i) => if i == -1 { None } else { Some(i) }
    RefArray(i) | RefNullArray(i) => if i == -1 { None } else { Some(i) }
    _ => None
  }
  guard ref_idx is Some(ref_idx) else { return }
  // Forward reference: ref_idx > type_idx
  if ref_idx > type_idx {
    // Must be in the same rec group
    if type_idx < type_rec_groups.length() && ref_idx < type_rec_groups.length() {
      if type_rec_groups[type_idx] != type_rec_groups[ref_idx] {
        raise UnknownType(ref_idx)
      }
    } else {
      // Type index out of bounds - this is an unknown type
      raise UnknownType(ref_idx)
    }
  }
}