// Projection pruning: after binding, compute which columns of each FROM
// table the query actually references, drop the rest from the pipeline
// layout, and rewrite every column reference to the narrowed indexes.
// Base tables materialize all their columns, so dropping a column is a
// zero-copy subset of each batch; joins, aggregation and projection
// then run over narrower rows. SELECT * references everything and
// degenerates to an identity rewrite.

///|
priv struct PrunePlan {
  bq : BoundQuery
  table_cols : Array[Array[Int32]] // kept columns per FROM table, new order
}

///|
fn prune_bound(
  bq : BoundQuery,
  schemas : Array[@catalog.TableSchema],
) -> PrunePlan {
  let ntab = schemas.length()
  let widths : Array[Int32] = []
  for s in schemas {
    widths.push(s.columns.length())
  }
  // referenced[t][local] over the pre-prune layout
  let needed : Array[Array[Bool]] = []
  for t in 0.. collect_refs(f, fn(i) { (step.table, i) }, needed)
      None => ()
    }
  }
  // flat-space expressions resolve through the join layout
  let total = bq.steps[bq.steps.length() - 1].right_offset +
    bq.steps[bq.steps.length() - 1].right_width
  let slot_of : Array[Int32] = Array::make(total, 0)
  let local_of : Array[Int32] = Array::make(total, 0)
  for step in bq.steps {
    for l in 0.. Unit {
    collect_refs(e, fn(i) { (slot_of[i], local_of[i]) }, needed)
  }
  for step in bq.steps {
    for e in step.keys_l {
      mark_flat(e)
    }
    match step.probe_filter {
      Some(f) => mark_flat(f)
      None => ()
    }
    match step.residual {
      Some(f) => mark_flat(f)
      None => ()
    }
  }
  match bq.deferred {
    Some(f) => mark_flat(f)
    None => ()
  }
  for e in bq.group_exprs {
    mark_flat(e)
  }
  for spec in bq.aggs {
    match spec.input {
      Some(e) => mark_flat(e)
      None => ()
    }
  }
  if !bq.aggregate {
    // AggOut-space projections reference the post-agg row, whose base
    // columns are already covered by group_exprs/aggs above
    for e in bq.out_exprs {
      mark_flat(e)
    }
    // ORDER BY keys sort the pre-projection row: their columns are
    // needed even when nothing else references them
    for k in bq.order_by {
      mark_flat(k.expr)
    }
    // window partition/order/input keys read the pre-projection row the
    // same way (their WinRef outputs mark no base columns)
    for w in bq.windows {
      match w.spec.input {
        Some(e) => mark_flat(e)
        None => ()
      }
      for e in w.partition {
        mark_flat(e)
      }
      for k in w.order {
        mark_flat(k.0)
      }
    }
  }
  // fast path: every column of every table is referenced
  let mut all_needed = true
  for t in 0.. new-flat map
  let new_off : Array[Int32] = Array::make(ntab, 0)
  let flat_map : Array[Int32] = Array::make(total, -1)
  let mut acc_w = 0
  for step in bq.steps {
    let t = step.table
    new_off[t] = acc_w
    for l in 0..= 0 {
        flat_map[step.right_offset + l] = acc_w + new_local[t][l]
      }
    }
    acc_w += new_width[t]
  }
  // per-table kept dtypes, then rebuilt steps with cumulative layouts
  let slot_dtypes : Array[Array[@types.DataType]] = []
  for t in 0.. Some(rewrite_refs(e, fn(i) { flat_map[i] }))
        None => None
      },
      out_type: spec.out_type,
      sep: spec.sep,
    })
  }
  let out_exprs = if bq.aggregate {
    bq.out_exprs
  } else {
    rewrite_all(bq.out_exprs, fn(i) { flat_map[i] })
  }
  let order_by = if bq.aggregate {
    bq.order_by
  } else {
    let out : Array[OrderSpec] = []
    for k in bq.order_by {
      out.push({
        expr: rewrite_refs(k.expr, fn(i) { flat_map[i] }),
        desc: k.desc,
        nulls_first: k.nulls_first,
        out_col: k.out_col,
      })
    }
    out
  }
  let windows = if bq.aggregate {
    bq.windows
  } else {
    let out : Array[BoundWin] = []
    for w in bq.windows {
      out.push({
        kind: w.kind,
        spec: {
          kind: w.spec.kind,
          input: match w.spec.input {
            Some(e) => Some(rewrite_refs(e, fn(i) { flat_map[i] }))
            None => None
          },
          out_type: w.spec.out_type,
          sep: w.spec.sep,
        },
        partition: rewrite_all(w.partition, fn(i) { flat_map[i] }),
        order: rewrite_order_exprs(w.order, fn(i) { flat_map[i] }),
      })
    }
    out
  }
  {
    bq: {
      steps,
      deferred,
      aggregate: bq.aggregate,
      group_exprs,
      group_types: bq.group_types,
      aggs,
      having: bq.having,
      out_exprs,
      windows,
      names: bq.names,
      out_types: bq.out_types,
      order_by,
      limit: bq.limit,
      offset: bq.offset,
      distinct: bq.distinct,
    },
    table_cols,
  }
}

///|
fn copy_dtypes(dts : Array[@types.DataType]) -> Array[@types.DataType] {
  let out : Array[@types.DataType] = []
  for dt in dts {
    out.push(dt)
  }
  out
}

///|
/// Mark every column a PhysExpr references. `resolve` maps a ColRef
/// index to (table, local column).
fn collect_refs(
  e : PhysExpr,
  resolve : (Int32) -> (Int32, Int32),
  needed : Array[Array[Bool]],
) -> Unit {
  match e {
    ColRef(i, _) => {
      let (t, l) = resolve(i)
      needed[t][l] = true
    }
    Const(_) => ()
    Promote(inner, _) => collect_refs(inner, resolve, needed)
    ArithE(_, l, r, _) => {
      collect_refs(l, resolve, needed)
      collect_refs(r, resolve, needed)
    }
    CmpE(_, l, r) => {
      collect_refs(l, resolve, needed)
      collect_refs(r, resolve, needed)
    }
    AndE(l, r) => {
      collect_refs(l, resolve, needed)
      collect_refs(r, resolve, needed)
    }
    OrE(l, r) => {
      collect_refs(l, resolve, needed)
      collect_refs(r, resolve, needed)
    }
    NotE(inner) => collect_refs(inner, resolve, needed)
    LikeE(inner, _) => collect_refs(inner, resolve, needed)
    IsNullE(inner, _) => collect_refs(inner, resolve, needed)
    ExtractE(inner, _) => collect_refs(inner, resolve, needed)
    InSetE(inner, _, _, _) => collect_refs(inner, resolve, needed)
    FuncE(_, args, _) =>
      for a in args {
        collect_refs(a, resolve, needed)
      }
    CastE(inner, _) => collect_refs(inner, resolve, needed)
    DateArithE(inner, _, _) => collect_refs(inner, resolve, needed)
    // a window output marker marks no base columns; the window's own
    // partition/order/input expressions are marked separately
    WinRef(_, _) => ()
    CaseE(whens, else_, _) => {
      for w in whens {
        collect_refs(w.0, resolve, needed)
        collect_refs(w.1, resolve, needed)
      }
      match else_ {
        Some(e) => collect_refs(e, resolve, needed)
        None => ()
      }
    }
  }
}

///|
/// Deep-copy an expression list with every ColRef index remapped.
fn rewrite_all(
  exprs : Array[PhysExpr],
  m : (Int32) -> Int32,
) -> Array[PhysExpr] {
  let out : Array[PhysExpr] = []
  for e in exprs {
    out.push(rewrite_refs(e, m))
  }
  out
}

///|
fn rewrite_opt(e : PhysExpr?, m : (Int32) -> Int32) -> PhysExpr? {
  match e {
    Some(inner) => Some(rewrite_refs(inner, m))
    None => None
  }
}

///|
/// Deep-copy ORDER BY keys with every ColRef index remapped.
fn rewrite_order_exprs(
  keys : Array[(PhysExpr, Bool, Bool?)],
  m : (Int32) -> Int32,
) -> Array[(PhysExpr, Bool, Bool?)] {
  let out : Array[(PhysExpr, Bool, Bool?)] = []
  for k in keys {
    out.push((rewrite_refs(k.0, m), k.1, k.2))
  }
  out
}

///|
fn rewrite_refs(e : PhysExpr, m : (Int32) -> Int32) -> PhysExpr {
  match e {
    ColRef(i, dt) => ColRef(m(i), dt)
    Const(v) => Const(v)
    Promote(inner, dt) => Promote(rewrite_refs(inner, m), dt)
    ArithE(op, l, r, dt) =>
      ArithE(op, rewrite_refs(l, m), rewrite_refs(r, m), dt)
    CmpE(op, l, r) => CmpE(op, rewrite_refs(l, m), rewrite_refs(r, m))
    AndE(l, r) => AndE(rewrite_refs(l, m), rewrite_refs(r, m))
    OrE(l, r) => OrE(rewrite_refs(l, m), rewrite_refs(r, m))
    NotE(inner) => NotE(rewrite_refs(inner, m))
    LikeE(inner, pattern) => LikeE(rewrite_refs(inner, m), pattern)
    IsNullE(inner, negated) => IsNullE(rewrite_refs(inner, m), negated)
    ExtractE(inner, field) => ExtractE(rewrite_refs(inner, m), field)
    InSetE(inner, set, negated, has_null) =>
      InSetE(rewrite_refs(inner, m), set, negated, has_null)
    FuncE(name, args, dt) => {
      let new_args : Array[PhysExpr] = []
      for a in args {
        new_args.push(rewrite_refs(a, m))
      }
      FuncE(name, new_args, dt)
    }
    CastE(inner, dt) => CastE(rewrite_refs(inner, m), dt)
    DateArithE(inner, qty, unit) =>
      DateArithE(rewrite_refs(inner, m), qty, unit)
    // window output markers survive pruning untouched: they index the
    // appended window columns, not the flat layout
    WinRef(i, dt) => WinRef(i, dt)
    CaseE(whens, else_, dt) => {
      let new_whens : Array[(PhysExpr, PhysExpr)] = []
      for w in whens {
        new_whens.push((rewrite_refs(w.0, m), rewrite_refs(w.1, m)))
      }
      let new_else = match else_ {
        Some(e) => Some(rewrite_refs(e, m))
        None => None
      }
      CaseE(new_whens, new_else, dt)
    }
  }
}