// cfg.ce — control flow graph
//
// Usage:
//   cell cfg --fn <N|name> <file>         Text CFG for function
//   cell cfg --dot --fn <N|name> <file>   DOT output for graphviz
//   cell cfg <file>                       Text CFG for all functions

var shop = use("internal/shop")

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 text(v)
  if (is_logical(v)) return v ? "true" : "false"
  return text(v)
}

var is_jump_op = function(op) {
  return op == "jump" || op == "jump_true" || op == "jump_false" || op == "jump_null" || op == "jump_not_null"
}

var is_conditional_jump = function(op) {
  return op == "jump_true" || op == "jump_false" || op == "jump_null" || op == "jump_not_null"
}

var is_terminator = function(op) {
  return op == "return" || op == "disrupt" || op == "tail_invoke" || op == "goinvoke"
}

var run = function() {
  var filename = null
  var fn_filter = null
  var show_dot = false
  var use_optimized = false
  var i = 0
  var compiled = null
  var main_name = null
  var fi = 0
  var func = null
  var fname = null

  while (i < length(args)) {
    if (args[i] == '--fn') {
      i = i + 1
      fn_filter = args[i]
    } else if (args[i] == '--dot') {
      show_dot = true
    } else if (args[i] == '--optimized') {
      use_optimized = true
    } else if (args[i] == '--help' || args[i] == '-h') {
      log.console("Usage: cell cfg [--fn <N|name>] [--dot] [--optimized] <file>")
      log.console("")
      log.console("  --fn <N|name>   Filter to function by index or name")
      log.console("  --dot           Output DOT format for graphviz")
      log.console("  --optimized     Use optimized IR")
      return null
    } else if (!starts_with(args[i], '-')) {
      filename = args[i]
    }
    i = i + 1
  }

  if (!filename) {
    log.console("Usage: cell cfg [--fn <N|name>] [--dot] [--optimized] <file>")
    return null
  }

  if (use_optimized) {
    compiled = shop.compile_file(filename)
  } else {
    compiled = shop.mcode_file(filename)
  }

  var fn_matches = function(index, name) {
    var match = null
    if (fn_filter == null) return true
    if (index >= 0 && fn_filter == text(index)) return true
    if (name != null) {
      match = search(name, fn_filter)
      if (match != null && match >= 0) return true
    }
    return false
  }

  var build_cfg = function(func) {
    var instrs = func.instructions
    var blocks = []
    var label_to_block = {}
    var pc_to_block = {}
    var label_to_pc = {}
    var block_start_pcs = {}
    var after_terminator = false
    var current_block = null
    var current_label = null
    var pc = 0
    var ii = 0
    var bi = 0
    var instr = null
    var op = null
    var n = 0
    var line_num = null
    var blk = null
    var last_instr_data = null
    var last_op = null
    var target_label = null
    var target_bi = null
    var edge_type = null

    if (instrs == null || length(instrs) == 0) return []

    // Pass 1: identify block start PCs
    block_start_pcs["0"] = true
    pc = 0
    ii = 0
    while (ii < length(instrs)) {
      instr = instrs[ii]
      if (is_array(instr)) {
        op = instr[0]
        if (after_terminator) {
          block_start_pcs[text(pc)] = true
          after_terminator = false
        }
        if (is_jump_op(op) || is_terminator(op)) {
          after_terminator = true
        }
        pc = pc + 1
      }
      ii = ii + 1
    }

    // Pass 2: map labels to PCs and mark as block starts
    pc = 0
    ii = 0
    while (ii < length(instrs)) {
      instr = instrs[ii]
      if (is_text(instr) && !starts_with(instr, "_nop_")) {
        label_to_pc[instr] = pc
        block_start_pcs[text(pc)] = true
      } else if (is_array(instr)) {
        pc = pc + 1
      }
      ii = ii + 1
    }

    // Pass 3: build basic blocks
    pc = 0
    ii = 0
    current_label = null
    while (ii < length(instrs)) {
      instr = instrs[ii]
      if (is_text(instr)) {
        if (!starts_with(instr, "_nop_")) {
          current_label = instr
        }
        ii = ii + 1
        continue
      }

      if (is_array(instr)) {
        if (block_start_pcs[text(pc)]) {
          if (current_block != null) {
            push(blocks, current_block)
          }
          current_block = {
            id: length(blocks),
            label: current_label,
            start_pc: pc,
            end_pc: pc,
            instrs: [],
            edges: [],
            first_line: null,
            last_line: null
          }
          current_label = null
        }

        if (current_block != null) {
          push(current_block.instrs, {pc: pc, instr: instr})
          current_block.end_pc = pc
          n = length(instr)
          line_num = instr[n - 2]
          if (line_num != null) {
            if (current_block.first_line == null) {
              current_block.first_line = line_num
            }
            current_block.last_line = line_num
          }
        }
        pc = pc + 1
      }
      ii = ii + 1
    }
    if (current_block != null) {
      push(blocks, current_block)
    }

    // Build block index
    bi = 0
    while (bi < length(blocks)) {
      pc_to_block[text(blocks[bi].start_pc)] = bi
      if (blocks[bi].label != null) {
        label_to_block[blocks[bi].label] = bi
      }
      bi = bi + 1
    }

    // Pass 4: compute edges
    bi = 0
    while (bi < length(blocks)) {
      blk = blocks[bi]
      if (length(blk.instrs) > 0) {
        last_instr_data = blk.instrs[length(blk.instrs) - 1]
        last_op = last_instr_data.instr[0]
        n = length(last_instr_data.instr)

        if (is_jump_op(last_op)) {
          if (last_op == "jump") {
            target_label = last_instr_data.instr[1]
          } else {
            target_label = last_instr_data.instr[2]
          }

          target_bi = label_to_block[target_label]
          if (target_bi != null) {
            edge_type = "jump"
            if (target_bi <= bi) {
              edge_type = "loop back-edge"
            }
            push(blk.edges, {target: target_bi, kind: edge_type})
          }

          if (is_conditional_jump(last_op)) {
            if (bi + 1 < length(blocks)) {
              push(blk.edges, {target: bi + 1, kind: "fallthrough"})
            }
          }
        } else if (is_terminator(last_op)) {
          push(blk.edges, {target: -1, kind: "EXIT (" + last_op + ")"})
        } else {
          if (bi + 1 < length(blocks)) {
            push(blk.edges, {target: bi + 1, kind: "fallthrough"})
          }
        }
      }
      bi = bi + 1
    }

    return blocks
  }

  var print_cfg_text = function(blocks, name) {
    var bi = 0
    var blk = null
    var header = null
    var ii = 0
    var idata = null
    var instr = null
    var op = null
    var n = 0
    var parts = null
    var j = 0
    var operands = null
    var ei = 0
    var edge = null
    var target_label = null

    log.compile(`\n=== ${name} ===`)

    if (length(blocks) == 0) {
      log.compile("  (empty)")
      return null
    }

    bi = 0
    while (bi < length(blocks)) {
      blk = blocks[bi]
      header = `  B${text(bi)}`
      if (blk.label != null) {
        header = header + ` "${blk.label}"`
      }
      header = header + ` [pc ${text(blk.start_pc)}-${text(blk.end_pc)}`
      if (blk.first_line != null) {
        if (blk.first_line == blk.last_line) {
          header = header + `, line ${text(blk.first_line)}`
        } else {
          header = header + `, lines ${text(blk.first_line)}-${text(blk.last_line)}`
        }
      }
      header = header + "]:"

      log.compile(header)

      ii = 0
      while (ii < length(blk.instrs)) {
        idata = blk.instrs[ii]
        instr = idata.instr
        op = instr[0]
        n = length(instr)
        parts = []
        j = 1
        while (j < n - 2) {
          push(parts, fmt_val(instr[j]))
          j = j + 1
        }
        operands = text(parts, ", ")
        log.compile(`    ${pad_right(text(idata.pc), 6)}${pad_right(op, 15)}${operands}`)
        ii = ii + 1
      }

      ei = 0
      while (ei < length(blk.edges)) {
        edge = blk.edges[ei]
        if (edge.target == -1) {
          log.compile(`    -> ${edge.kind}`)
        } else {
          target_label = blocks[edge.target].label
          if (target_label != null) {
            log.compile(`    -> B${text(edge.target)} "${target_label}" (${edge.kind})`)
          } else {
            log.compile(`    -> B${text(edge.target)} (${edge.kind})`)
          }
        }
        ei = ei + 1
      }

      log.compile("")
      bi = bi + 1
    }
    return null
  }

  var print_cfg_dot = function(blocks, name) {
    var safe_name = replace(replace(name, '"', '\\"'), ' ', '_')
    var bi = 0
    var blk = null
    var label_text = null
    var ii = 0
    var idata = null
    var instr = null
    var op = null
    var n = 0
    var parts = null
    var j = 0
    var operands = null
    var ei = 0
    var edge = null
    var style = null

    log.compile(`digraph "${safe_name}" {`)
    log.compile("  rankdir=TB;")
    log.compile("  node [shape=record, fontname=monospace, fontsize=10];")

    bi = 0
    while (bi < length(blocks)) {
      blk = blocks[bi]
      label_text = "B" + text(bi)
      if (blk.label != null) {
        label_text = label_text + " (" + blk.label + ")"
      }
      label_text = label_text + "\\npc " + text(blk.start_pc) + "-" + text(blk.end_pc)
      if (blk.first_line != null) {
        label_text = label_text + "\\nline " + text(blk.first_line)
      }
      label_text = label_text + "|"

      ii = 0
      while (ii < length(blk.instrs)) {
        idata = blk.instrs[ii]
        instr = idata.instr
        op = instr[0]
        n = length(instr)
        parts = []
        j = 1
        while (j < n - 2) {
          push(parts, fmt_val(instr[j]))
          j = j + 1
        }
        operands = text(parts, ", ")
        label_text = label_text + text(idata.pc) + " " + op + " " + replace(operands, '"', '\\"') + "\\l"
        ii = ii + 1
      }

      log.compile("  B" + text(bi) + " [label=\"{" + label_text + "}\"];")
      bi = bi + 1
    }

    // Edges
    bi = 0
    while (bi < length(blocks)) {
      blk = blocks[bi]
      ei = 0
      while (ei < length(blk.edges)) {
        edge = blk.edges[ei]
        if (edge.target >= 0) {
          style = ""
          if (edge.kind == "loop back-edge") {
            style = " [style=bold, color=red, label=\"loop\"]"
          } else if (edge.kind == "fallthrough") {
            style = " [style=dashed]"
          }
          log.compile(`  B${text(bi)} -> B${text(edge.target)}${style};`)
        }
        ei = ei + 1
      }
      bi = bi + 1
    }

    log.compile("}")
    return null
  }

  var process_function = function(func, name, index) {
    var blocks = build_cfg(func)
    if (show_dot) {
      print_cfg_dot(blocks, name)
    } else {
      print_cfg_text(blocks, name)
    }
    return null
  }

  // Process functions
  main_name = compiled.name != null ? compiled.name : "<main>"

  if (compiled.main != null) {
    if (fn_matches(-1, main_name)) {
      process_function(compiled.main, main_name, -1)
    }
  }

  if (compiled.functions != null) {
    fi = 0
    while (fi < length(compiled.functions)) {
      func = compiled.functions[fi]
      fname = func.name != null ? func.name : "<anonymous>"
      if (fn_matches(fi, fname)) {
        process_function(func, `[${text(fi)}] ${fname}`, fi)
      }
      fi = fi + 1
    }
  }

  return null
}

run()
$stop()