// ir_stats.cm — IR statistics, fingerprinting, and canonical printing
//
// Usage: var ir_stats = use("ir_stats")
//   ir_stats.detailed_stats(func)
//   ir_stats.ir_fingerprint(func)
//   ir_stats.canonical_ir(func, name, opts)
//   ir_stats.type_snapshot(slot_types)
//   ir_stats.type_delta(before_types, after_types)

var json = use("json")

// --- Category maps ---

var load_ops = {
  load_field: true, load_index: true, load_dynamic: true,
  get: true
}
var store_ops = {
  store_field: true, store_index: true, store_dynamic: true,
  set_var: true, put: true, push: true
}
var branch_ops = {
  jump: true, jump_true: true, jump_false: true, jump_not_null: true
}
var call_ops = {
  invoke: true, goinvoke: true
}
var guard_ops = {
  is_int: true, is_text: true, is_num: true, is_bool: true,
  is_null: true, is_array: true, is_func: true, is_record: true,
  is_stone: true
}
var arith_ops = {
  add_int: true, sub_int: true, mul_int: true, div_int: true, mod_int: true,
  add_float: true, sub_float: true, mul_float: true, div_float: true, mod_float: true,
  concat: true, neg_int: true, neg_float: true,
  bitnot: true, bitand: true, bitor: true, bitxor: true,
  shl: true, shr: true, ushr: true
}
var move_ops = {
  move: true
}
var const_ops = {
  int: true, true: true, false: true, null: true
}

var nop_reasons = {
  tc: "tc",
  bl: "bl",
  mv: "mv",
  dj: "dj",
  ur: "ur"
}

var category_tag = function(op) {
  if (guard_ops[op] == true) { return "guard" }
  if (branch_ops[op] == true) { return "branch" }
  if (load_ops[op] == true) { return "load" }
  if (store_ops[op] == true) { return "store" }
  if (call_ops[op] == true) { return "call" }
  if (arith_ops[op] == true) { return "arith" }
  if (move_ops[op] == true) { return "move" }
  if (const_ops[op] == true) { return "const" }
  return null
}

// --- detailed_stats ---

var detailed_stats = function(func) {
  var instructions = func.instructions
  var stats = {
    instr: 0, nop: 0,
    load: 0, store: 0, branch: 0, call: 0,
    guard: 0, arith: 0, move: 0, const: 0,
    label: 0, other: 0
  }
  var i = 0
  var instr = null
  var op = null
  var num = 0

  if (instructions == null) {
    return stats
  }

  num = length(instructions)
  while (i < num) {
    instr = instructions[i]
    if (is_text(instr)) {
      if (starts_with(instr, "_nop_")) {
        stats.nop = stats.nop + 1
        stats.instr = stats.instr + 1
      } else {
        stats.label = stats.label + 1
      }
    } else if (is_array(instr)) {
      stats.instr = stats.instr + 1
      op = instr[0]
      if (op == "access" && !is_number(instr[2]) && !is_logical(instr[2])) {
        stats.load = stats.load + 1
      } else if (op == "access") {
        stats.const = stats.const + 1
      } else if (load_ops[op] == true) {
        stats.load = stats.load + 1
      } else if (store_ops[op] == true) {
        stats.store = stats.store + 1
      } else if (branch_ops[op] == true) {
        stats.branch = stats.branch + 1
      } else if (call_ops[op] == true) {
        stats.call = stats.call + 1
      } else if (guard_ops[op] == true) {
        stats.guard = stats.guard + 1
      } else if (arith_ops[op] == true) {
        stats.arith = stats.arith + 1
      } else if (move_ops[op] == true) {
        stats.move = stats.move + 1
      } else if (const_ops[op] == true) {
        stats.const = stats.const + 1
      } else {
        stats.other = stats.other + 1
      }
    }
    i = i + 1
  }
  return stats
}

// --- ir_fingerprint ---
// djb2 hash computed over the JSON-encoded instructions

var djb2 = function(s) {
  var chars = array(s)
  var hash = 5381
  var i = 0
  var num = length(chars)
  while (i < num) {
    hash = ((hash * 33) + number(chars[i])) % 4294967296
    i = i + 1
  }
  return text(hash, 16)
}

var ir_fingerprint = function(func) {
  return djb2(json.encode(func.instructions))
}

// --- canonical_ir ---

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

var nop_reason = function(s) {
  // extract reason from _nop_XX_NNN
  var parts = array(s, "_")
  // parts: ["", "nop", "XX", "NNN"]
  if (length(parts) >= 3) {
    return parts[2]
  }
  return "?"
}

var fmt_operand = function(v) {
  if (is_null(v)) {
    return "null"
  }
  if (is_number(v)) {
    return text(v)
  }
  if (is_text(v)) {
    return `"${v}"`
  }
  if (is_logical(v)) {
    if (v) { return "true" }
    return "false"
  }
  return text(v)
}

var canonical_ir = function(func, name, opts) {
  var instructions = func.instructions
  var nr_args = func.nr_args != null ? func.nr_args : 0
  var nr_slots = func.nr_slots != null ? func.nr_slots : 0
  var show_nops = false
  var show_types = false
  var slot_types = null
  var lines = []
  var i = 0
  var instr = null
  var op = null
  var n = 0
  var parts = null
  var j = 0
  var idx_str = null
  var op_str = null
  var operands = null
  var suffix = null
  var tag = null
  var typ = null
  var reason = null
  var num = 0

  if (opts != null) {
    if (opts.show_nops == true) { show_nops = true }
    if (opts.show_types == true) { show_types = true }
    if (opts.slot_types != null) { slot_types = opts.slot_types }
  }

  lines[] = `fn ${name != null ? name : "<anon>"} (args=${text(nr_args)}, slots=${text(nr_slots)})`

  if (instructions == null) {
    return text(lines, "\n")
  }

  num = length(instructions)
  while (i < num) {
    instr = instructions[i]

    if (is_text(instr)) {
      if (starts_with(instr, "_nop_")) {
        if (show_nops) {
          reason = nop_reason(instr)
          idx_str = pad_right(`@${text(i)}`, 6)
          lines[] = `  ${idx_str}--- nop (${reason}) ---`
        }
      } else {
        lines[] = `      ${instr}:`
      }
      i = i + 1
      continue
    }

    if (!is_array(instr)) {
      i = i + 1
      continue
    }

    op = instr[0]
    n = length(instr)
    parts = []
    j = 1
    while (j < n - 2) {
      if (is_number(instr[j]) && op != "int" && !(op == "access" && j == 2)) {
        parts[] = `s${text(instr[j])}`
      } else {
        parts[] = fmt_operand(instr[j])
      }
      j = j + 1
    }
    operands = text(parts, ", ")

    idx_str = pad_right(`@${text(i)}`, 6)
    op_str = pad_right(op, 16)
    suffix = ""

    tag = category_tag(op)

    if (show_types && slot_types != null) {
      // show type for dest slot if known
      if (is_number(instr[1])) {
        typ = slot_types[text(instr[1])]
        if (typ != null) {
          suffix = `; -> ${typ}`
        }
      }
      if (tag != null) {
        suffix = suffix + `  [${tag}]`
      }
    } else if (tag != null) {
      suffix = suffix + `; [${tag}]`
    }

    if (length(suffix) > 0) {
      lines[] = `  ${idx_str}${op_str}${operands} ${suffix}`
    } else {
      lines[] = `  ${idx_str}${op_str}${operands}`
    }

    i = i + 1
  }

  return text(lines, "\n")
}

// --- type_snapshot ---

var type_snapshot = function(slot_types) {
  if (slot_types == null) {
    return {}
  }
  return stone(record(slot_types))
}

// --- type_delta ---

var type_delta = function(before_types, after_types) {
  var result = {
    added: {},
    removed: {},
    strengthened: {},
    weakened: {}
  }
  var bt = before_types != null ? before_types : {}
  var at = after_types != null ? after_types : {}
  var keys = null
  var i = 0
  var k = null
  var bv = null
  var av = null

  // check after for added/changed
  keys = array(at)
  i = 0
  while (i < length(keys)) {
    k = keys[i]
    av = at[k]
    bv = bt[k]
    if (bv == null) {
      result.added[k] = av
    } else if (bv != av) {
      if (bv == "unknown" || (bv == "num" && (av == "int" || av == "float"))) {
        result.strengthened[k] = {from: bv, to: av}
      } else if (av == "unknown" || (av == "num" && (bv == "int" || bv == "float"))) {
        result.weakened[k] = {from: bv, to: av}
      } else {
        result.strengthened[k] = {from: bv, to: av}
      }
    }
    i = i + 1
  }

  // check before for removed
  keys = array(bt)
  i = 0
  while (i < length(keys)) {
    k = keys[i]
    if (at[k] == null) {
      result.removed[k] = bt[k]
    }
    i = i + 1
  }

  return result
}

return {
  detailed_stats: detailed_stats,
  ir_fingerprint: ir_fingerprint,
  canonical_ir: canonical_ir,
  type_snapshot: type_snapshot,
  type_delta: type_delta,
  category_tag: category_tag
}
