// dump_types.cm — show streamlined IR with type annotations
//
// Usage: ./cell --core . dump_types.cm <file.ce|file.cm>

var fd = use("fd")
var json = use("json")
var tokenize = use("tokenize")
var parse = use("parse")
var fold = use("fold")
var mcode = use("mcode")
var streamline = use("streamline")

if (length(args) < 1) {
  print("usage: cell --core . dump_types.cm <file>")
  return
}

var filename = args[0]
var src = text(fd.slurp(filename))
var tok = tokenize(src, filename)
var ast = parse(tok.tokens, src, filename, tokenize)
var folded = fold(ast)
var compiled = mcode(folded)
var optimized = streamline(compiled)

// Type constants
def T_UNKNOWN = "unknown"
def T_INT = "int"
def T_FLOAT = "float"
def T_NUM = "num"
def T_TEXT = "text"
def T_BOOL = "bool"
def T_NULL = "null"
def T_ARRAY = "array"
def T_RECORD = "record"
def T_FUNCTION = "function"

def int_result_ops = {
  bitnot: true, bitand: true, bitor: true,
  bitxor: true, shl: true, shr: true, ushr: true
}
def bool_result_ops = {
  eq_int: true, ne_int: true, lt_int: true, gt_int: true,
  le_int: true, ge_int: true,
  eq_float: true, ne_float: true, lt_float: true, gt_float: true,
  le_float: true, ge_float: true,
  eq_text: true, ne_text: true, lt_text: true, gt_text: true,
  le_text: true, ge_text: true,
  eq_bool: true, ne_bool: true,
  not: true, and: true, or: true,
  is_int: true, is_text: true, is_num: true,
  is_bool: true, is_null: true, is_identical: true,
  is_array: true, is_func: true, is_record: true, is_stone: true
}

var access_value_type = function(val) {
  if (is_number(val)) {
    return is_integer(val) ? T_INT : T_FLOAT
  }
  if (is_text(val)) {
    return T_TEXT
  }
  return T_UNKNOWN
}

var track_types = function(slot_types, instr) {
  var op = instr[0]
  var src_type = null
  if (op == "access") {
    slot_types[text(instr[1])] = access_value_type(instr[2])
  } else if (op == "int") {
    slot_types[text(instr[1])] = T_INT
  } else if (op == "true" || op == "false") {
    slot_types[text(instr[1])] = T_BOOL
  } else if (op == "null") {
    slot_types[text(instr[1])] = T_NULL
  } else if (op == "move") {
    src_type = slot_types[text(instr[2])]
    slot_types[text(instr[1])] = src_type != null ? src_type : T_UNKNOWN
  } else if (int_result_ops[op] == true) {
    slot_types[text(instr[1])] = T_INT
  } else if (op == "concat") {
    slot_types[text(instr[1])] = T_TEXT
  } else if (bool_result_ops[op] == true) {
    slot_types[text(instr[1])] = T_BOOL
  } else if (op == "typeof") {
    slot_types[text(instr[1])] = T_TEXT
  } else if (op == "array") {
    slot_types[text(instr[1])] = T_ARRAY
  } else if (op == "record") {
    slot_types[text(instr[1])] = T_RECORD
  } else if (op == "function") {
    slot_types[text(instr[1])] = T_FUNCTION
  } else if (op == "invoke" || op == "tail_invoke") {
    slot_types[text(instr[2])] = T_UNKNOWN
  } else if (op == "load_field" || op == "load_index" || op == "load_dynamic") {
    slot_types[text(instr[1])] = T_UNKNOWN
  } else if (op == "pop" || op == "get") {
    slot_types[text(instr[1])] = T_UNKNOWN
  } else if (op == "length") {
    slot_types[text(instr[1])] = T_INT
  } else if (op == "add" || op == "subtract" || op == "multiply" ||
             op == "divide" || op == "modulo" || op == "pow" || op == "negate") {
    slot_types[text(instr[1])] = T_UNKNOWN
  }
  return null
}

var pad_right = function(s, w) {
  var r = s
  while (length(r) < w) {
    r = r + " "
  }
  return r
}

var fmt_val = function(v) {
  if (is_null(v)) {
    return "null"
  }
  if (is_number(v)) {
    return text(v)
  }
  if (is_text(v)) {
    return `"${v}"`
  }
  if (is_object(v)) {
    return json.encode(v)
  }
  if (is_logical(v)) {
    return v ? "true" : "false"
  }
  return text(v)
}

// Build type annotation string for an instruction
var type_annotation = function(slot_types, instr) {
  var n = length(instr)
  var parts = []
  var j = 1
  var v = null
  var t = null
  while (j < n - 2) {
    v = instr[j]
    if (is_number(v)) {
      t = slot_types[text(v)]
      if (t != null && t != T_UNKNOWN) {
        push(parts, `s${text(v)}:${t}`)
      }
    }
    j = j + 1
  }
  if (length(parts) == 0) {
    return ""
  }
  return text(parts, " ")
}

var dump_function_typed = function(func, name) {
  var nr_args = func.nr_args != null ? func.nr_args : 0
  var nr_slots = func.nr_slots != null ? func.nr_slots : 0
  var instrs = func.instructions
  var slot_types = {}
  var i = 0
  var pc = 0
  var instr = null
  var op = null
  var n = 0
  var annotation = null
  var operand_parts = null
  var j = 0
  var operands = null
  var pc_str = null
  var op_str = null
  var line = null
  print(`\n=== ${name} (args=${text(nr_args)}, slots=${text(nr_slots)}) ===`)
  if (instrs == null || length(instrs) == 0) {
    print("  (empty)")
    return null
  }
  while (i < length(instrs)) {
    instr = instrs[i]
    if (is_text(instr)) {
      if (starts_with(instr, "_nop_")) {
        i = i + 1
        continue
      }
      slot_types = {}
      print(`${instr}:`)
    } else if (is_array(instr)) {
      op = instr[0]
      n = length(instr)
      annotation = type_annotation(slot_types, instr)
      operand_parts = []
      j = 1
      while (j < n - 2) {
        push(operand_parts, fmt_val(instr[j]))
        j = j + 1
      }
      operands = text(operand_parts, ", ")
      pc_str = pad_right(text(pc), 5)
      op_str = pad_right(op, 14)
      line = pad_right(`  ${pc_str} ${op_str} ${operands}`, 50)
      if (length(annotation) > 0) {
        print(`${line} ; ${annotation}`)
      } else {
        print(line)
      }
      track_types(slot_types, instr)
      pc = pc + 1
    }
    i = i + 1
  }
  return null
}

var main_name = null
var fi = 0
var func = null
var fname = null

// Dump main
if (optimized.main != null) {
  main_name = optimized.name != null ? optimized.name : "<main>"
  dump_function_typed(optimized.main, main_name)
}

// Dump sub-functions
if (optimized.functions != null) {
  fi = 0
  while (fi < length(optimized.functions)) {
    func = optimized.functions[fi]
    fname = func.name != null ? func.name : `<func_${text(fi)}>`
    dump_function_typed(func, `[${text(fi)}] ${fname}`)
    fi = fi + 1
  }
}