voltaic-driver 0.2.0

use voltaic_ast::BinaryOp;

use voltaic_ast::Block;

use voltaic_ast::CaseArm;

use voltaic_ast::ClassMember;

use voltaic_ast::ClauseBody;

use voltaic_ast::Condition;

use voltaic_ast::ConditionOperand;

use voltaic_ast::Expr;

use voltaic_ast::FieldInit;

use voltaic_ast::Guard;

use voltaic_ast::Item;

use voltaic_ast::Program;

use voltaic_ast::Qualifier;

use voltaic_ast::Section;

use voltaic_ast::Stmt;

use voltaic_ast::StrInterpPart;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Pass {
  BlankNumbers,
  Pipelines,
}
#[allow(unused_imports)]
use Pass::*;

pub fn rewrite_program(program: &mut Program, step: Pass) {
  for item in program.items.iter_mut() {
    rewrite_item(item, step);
  }
  ();
}

pub fn rewrite_item(item: &mut Item, step: Pass) {
  match item {
    Item::Fn(f) => {
      for clause in f.clauses.iter_mut() {
        match &mut clause.body {
          ClauseBody::Plain(block) => rewrite_block(block, step),
          ClauseBody::Guarded(rows) => {
            for row in rows.iter_mut() {
              rewrite_guard(row, step);
            }
            ()
          }
        };
        for stmt in clause.wheres.iter_mut() {
          rewrite_stmt(stmt, step);
        }
      }
      ()
    }
    Item::Instance(block) => {
      for member in block.members.iter_mut() {
        if let ClassMember::Method(f) = member {
          for clause in f.clauses.iter_mut() {
            match &mut clause.body {
              ClauseBody::Plain(b) => rewrite_block(b, step),
              ClauseBody::Guarded(rows) => {
                for row in rows.iter_mut() {
                  rewrite_guard(row, step);
                }
                ()
              }
            };
            for stmt in clause.wheres.iter_mut() {
              rewrite_stmt(stmt, step);
            }
          }
        }
      }
      ()
    }
    Item::Test(t) => rewrite_block(&mut t.body, step),
    Item::Const(c) => rewrite_expr(&mut c.value, step),
    _ => (),
  };
}

pub fn rewrite_block(block: &mut Block, step: Pass) {
  for stmt in block.stmts.iter_mut() {
    rewrite_stmt(stmt, step);
  }
  ();
}

pub fn rewrite_stmt(stmt: &mut Stmt, step: Pass) {
  rewrite_expr(&mut stmt.expr, step);
  ();
}

pub fn rewrite_guard(row: &mut Guard, step: Pass) {
  if let Some(cond) = &mut row.condition {
    rewrite_condition(cond, step);
  }
  rewrite_block(&mut row.body, step);
  ();
}

pub fn rewrite_condition(cond: &mut Condition, step: Pass) {
  for operand in cond.operands.iter_mut() {
    match operand {
      ConditionOperand::Expr(e) => rewrite_expr(e, step),
      ConditionOperand::Pattern { source, .. } => rewrite_expr(source, step),
    };
  }
  ();
}

pub fn rewrite_arm(arm: &mut CaseArm, step: Pass) {
  if let Some(cond) = &mut arm.r#guard {
    rewrite_condition(cond, step);
  }
  rewrite_expr(&mut arm.body, step);
  ();
}

pub fn rewrite_field(field: &mut FieldInit, step: Pass) {
  if let Some(e) = &mut field.value {
    rewrite_expr(e, step);
  }
  ();
}

pub fn rewrite_part(part: &mut StrInterpPart, step: Pass) {
  match part {
    StrInterpPart::Interp { expr, .. } => rewrite_expr(expr, step),
    StrInterpPart::Text(_) => (),
  };
}

pub fn rewrite_qualifier(q: &mut Qualifier, step: Pass) {
  match q {
    Qualifier::Draw { source, .. } => rewrite_expr(source, step),
    Qualifier::Guard(e) => rewrite_expr(e, step),
  };
}

pub fn rewrite_section(s: &mut Section, step: Pass) {
  match s {
    Section::Right { operand, .. } => rewrite_expr(operand, step),
    Section::Left { operand, .. } => rewrite_expr(operand, step),
    Section::Field(_) => (),
  };
}

pub fn rewrite_expr(e: &mut Expr, step: Pass) {
  descend(e, step);
  apply_pass(e, step);
  ();
}

pub fn descend(e: &mut Expr, step: Pass) {
  match e {
    Expr::Call { callee, args, .. } => {
      rewrite_expr(callee, step);
      for a in args.iter_mut() {
        rewrite_expr(a, step);
      }
      ()
    }
    Expr::Field { recv, .. } => rewrite_expr(recv, step),
    Expr::Dot { recv, args, .. } => {
      rewrite_expr(recv, step);
      for a in args.iter_mut() {
        rewrite_expr(a, step);
      }
      ()
    }
    Expr::Index { recv, index, .. } => {
      rewrite_expr(recv, step);
      rewrite_expr(index, step)
    }
    Expr::Binary { lhs, rhs, .. } => {
      rewrite_expr(lhs, step);
      rewrite_expr(rhs, step)
    }
    Expr::UserBinary { lhs, rhs, .. } => {
      rewrite_expr(lhs, step);
      rewrite_expr(rhs, step)
    }
    Expr::Over {
      function, place, ..
    } => {
      rewrite_expr(function, step);
      rewrite_expr(place, step)
    }
    Expr::Unary { operand, .. } => rewrite_expr(operand, step),
    Expr::Paren { inner, .. } => rewrite_expr(inner, step),
    Expr::Lambda { body, .. } => rewrite_expr(body, step),
    Expr::AsyncBlock { body, .. } => rewrite_block(body, step),
    Expr::UnsafeBlock { body, .. } => rewrite_block(body, step),
    Expr::TryBlock { body, .. } => rewrite_block(body, step),
    Expr::BindBlock { body, .. } => rewrite_block(body, step),
    Expr::Ascribe { inner, .. } => rewrite_expr(inner, step),
    Expr::StrInterp { parts, .. } => {
      for p in parts.iter_mut() {
        rewrite_part(p, step);
      }
      ()
    }
    Expr::Tuple { items, .. } => {
      for i in items.iter_mut() {
        rewrite_expr(i, step);
      }
      ()
    }
    Expr::ArrayRepeat { value, count, .. } => {
      rewrite_expr(value, step);
      rewrite_expr(count, step)
    }
    Expr::Record { head, fields, .. } => {
      rewrite_expr(head, step);
      for f in fields.iter_mut() {
        rewrite_field(f, step);
      }
      ()
    }
    Expr::List { items, .. } => {
      for i in items.iter_mut() {
        rewrite_expr(i, step);
      }
      ()
    }
    Expr::Comprehension {
      element,
      qualifiers,
      ..
    } => {
      rewrite_expr(element, step);
      for q in qualifiers.iter_mut() {
        rewrite_qualifier(q, step);
      }
      ()
    }
    Expr::If {
      condition,
      then_block,
      r#otherwise,
      ..
    } => {
      rewrite_condition(condition, step);
      rewrite_block(then_block, step);
      if let Some(block) = r#otherwise {
        rewrite_block(block, step);
      }
      ()
    }
    Expr::Block { body, .. } => rewrite_block(body, step),
    Expr::Try { inner, .. } => rewrite_expr(inner, step),
    Expr::Section { section, .. } => rewrite_section(section, step),
    Expr::Case {
      scrutinee, arms, ..
    } => {
      rewrite_expr(scrutinee, step);
      for a in arms.iter_mut() {
        rewrite_arm(a, step);
      }
      ()
    }
    Expr::Ladder { rows, .. } => {
      for row in rows.iter_mut() {
        rewrite_guard(row, step);
      }
      ()
    }
    Expr::Binding(b) => {
      rewrite_expr(&mut b.value, step);
      if let Some(block) = &mut b.r#otherwise {
        rewrite_block(block, step);
      }
      ()
    }
    Expr::Item { item, .. } => rewrite_item(item, step),
    Expr::Return { value, .. } => {
      if let Some(v) = value {
        rewrite_expr(v, step);
      }
      ()
    }
    Expr::Draw { value, .. } => rewrite_expr(value, step),
    Expr::Await { value, .. } => rewrite_expr(value, step),
    Expr::Assign { target, value, .. } => {
      rewrite_expr(target, step);
      rewrite_expr(value, step)
    }
    Expr::Guard {
      source,
      r#otherwise,
      ..
    } => {
      rewrite_expr(source, step);
      rewrite_block(r#otherwise, step)
    }
    Expr::While {
      condition, body, ..
    } => {
      rewrite_condition(condition, step);
      rewrite_block(body, step)
    }
    Expr::For { source, body, .. } => {
      rewrite_expr(source, step);
      rewrite_block(body, step)
    }
    Expr::Loop { body, .. } => rewrite_block(body, step),
    Expr::Break { value, .. } => {
      if let Some(v) = value {
        rewrite_expr(v, step);
      }
      ()
    }
    Expr::MacroCall { args, .. } => {
      for a in args.iter_mut() {
        rewrite_expr(a, step);
      }
      ()
    }
    Expr::Path { .. }
    | Expr::MethodRef { .. }
    | Expr::StrLit { .. }
    | Expr::RawLit { .. }
    | Expr::Unit { .. }
    | Expr::Num { .. }
    | Expr::Bool { .. }
    | Expr::Char { .. }
    | Expr::SelfField { .. }
    | Expr::SelfType { .. }
    | Expr::Ident(_)
    | Expr::Continue { .. }
    | Expr::RustHatch { .. } => (),
  };
}

pub fn apply_pass(e: &mut Expr, pass: Pass) {
  match pass {
    Pass::BlankNumbers => blank_number(e),
    Pass::Pipelines => collapse_pipeline(e),
  };
}

pub(crate) fn blank_number(e: &mut Expr) {
  if let Expr::Num { text, .. } = e {
    *text = "9".to_string();
  }
  ();
}

pub(crate) fn collapse_pipeline(e: &mut Expr) {
  let Expr::Binary { op, lhs, rhs, span } = e else {
    return ();
  };
  if *op != BinaryOp::Pipeline {
    return ();
  }
  let mut named = None;
  let mut found: Vec<Expr> = Vec::new();
  if let Expr::MethodRef { name, .. } = rhs.as_ref() {
    named = Some(name.clone());
  }
  if let Expr::Call { callee, args, .. } = rhs.as_ref()
    && let Expr::MethodRef { name, .. } = callee.as_ref()
  {
    named = Some(name.clone());
    found = args.as_slice().to_vec();
  }
  let Some(name) = named else {
    return ();
  };
  *e = Expr::Dot {
    recv: lhs.clone(),
    name,
    args: found,
    span: *span,
  };
  ();
}

#[cfg(test)]
mod __vlt_tests {
  use super::*;

  #[test]
  fn a_leaf_has_no_children_so_descending_it_is_a_no_op() {
    let mut e = Expr::Bool {
      value: true,
      span: (0, 1),
    };
    let before = e.clone();
    rewrite_expr(&mut e, Pass::BlankNumbers);
    assert_eq!(e, before);
  }

  #[test]
  fn the_traversal_reaches_a_numeral_in_every_expression_position() {
    let source = "f : u64 -> u64\nf n =\n  m = n + 1\n  xs = [2, 3] : Vec<u64>\n  case m of\n    4 -> 5\n    _ -> m * 6\n";
    let mut program = voltaic_parser::parse_source(source).expect("should parse");
    rewrite_program(&mut program, Pass::BlankNumbers);
    let shown = format!("{:?}", program);
    for digit in &["\"1\"", "\"2\"", "\"3\"", "\"5\"", "\"6\""] {
      assert!(!shown.as_str().contains(*digit));
    }
  }

  #[test]
  fn a_numeral_in_a_pattern_is_left_alone_because_a_pass_must_not_change_what_matches() {
    let source = "f : u64 -> u64\nf n = case n of\n  4 -> 5\n  _ -> 6\n";
    let mut program = voltaic_parser::parse_source(source).expect("should parse");
    rewrite_program(&mut program, Pass::BlankNumbers);
    let shown = format!("{:?}", program);
    assert!(shown.as_str().contains("\"4\""));
    assert!(!shown.as_str().contains("\"5\""));
    assert!(!shown.as_str().contains("\"6\""));
  }

  #[test]
  fn a_numeral_the_traversal_cannot_reach_would_survive_which_is_what_this_pins() {
    let source = "f : u64 -> u64\nf n = n + 7\n";
    let mut program = voltaic_parser::parse_source(source).expect("should parse");
    let before = format!("{:?}", program);
    assert!(before.as_str().contains("\"7\""));
    rewrite_program(&mut program, Pass::BlankNumbers);
    let after = format!("{:?}", program);
    assert!(!after.as_str().contains("\"7\""));
  }

  #[test]
  fn a_pipeline_with_no_arguments_collapses_now_that_the_printer_writes_the_parentheses() {
    let source = "f : &str -> usize\nf s = s |> .len\n";
    let mut program = voltaic_parser::parse_source(source).expect("should parse");
    rewrite_program(&mut program, Pass::Pipelines);
    let shown = format!("{:?}", program);
    assert!(!shown.as_str().contains("Pipeline"));
    assert!(shown.as_str().contains("Dot"));
  }

  #[test]
  fn a_pipeline_with_arguments_keeps_them_receiver_first() {
    let source = "f : u64 -> u64\nf n = n |> .max 3\n";
    let mut program = voltaic_parser::parse_source(source).expect("should parse");
    rewrite_program(&mut program, Pass::Pipelines);
    let shown = format!("{:?}", program);
    assert!(shown.as_str().contains("Dot"));
    assert!(shown.as_str().contains("\"max\""));
    assert!(shown.as_str().contains("\"3\""));
  }

  #[test]
  fn a_chain_collapses_every_stage_arguments_or_not() {
    let source = "f : u64 -> u64\nf n = n |> .max 3 |> .min 8\n";
    let mut program = voltaic_parser::parse_source(source).expect("should parse");
    rewrite_program(&mut program, Pass::Pipelines);
    let shown = format!("{:?}", program);
    assert!(!shown.as_str().contains("Pipeline"));
  }

  #[test]
  fn a_mixed_chain_collapses_both_stages_the_one_with_arguments_and_the_one_without() {
    let source = "f : &str -> usize\nf s = s |> .replacen \"a\" \"b\" 1 |> .len\n";
    let mut program = voltaic_parser::parse_source(source).expect("should parse");
    rewrite_program(&mut program, Pass::Pipelines);
    let shown = format!("{:?}", program);
    assert!(shown.as_str().contains("Dot"));
    assert!(!shown.as_str().contains("Pipeline"));
  }

  #[test]
  fn a_binary_that_is_not_a_pipeline_is_left_alone() {
    let source = "f : u64 -> u64\nf n = n + 1\n";
    let mut program = voltaic_parser::parse_source(source).expect("should parse");
    let before = format!("{:?}", program);
    rewrite_program(&mut program, Pass::Pipelines);
    let after = format!("{:?}", program);
    assert_eq!(after, before);
  }
}