cmtc 0.1.2

use std::cell::RefCell;
use std::rc::Rc;

use cmtrs::*;
use gen_ir::ModuleMold;

// Counter interface, should be generated
// *************************************
struct Counter {
  // as generate context
  module: cmtrs::gen_ir::ModuleMold,
  // as instance
  ctx_view: Box<dyn InstanceView>,
  // params
  pub t: Type,
}

struct CounterItfc {
  // as generate context
  module: cmtrs::gen_ir::ModuleMold,
  // parameters and sub_itfc
  params: CounterParams,
}

struct CounterIO {
  pub set_val: Var,
  pub count: Var,
}

impl ItfcIO for CounterIO {}

struct CounterParams {
  pub t: Type,
}
impl ItfcParams for CounterParams {
  fn id_str(&self) -> String {
    format!("{}", self.t.name())
  }
}

impl Itfc for CounterItfc {
  type IO = CounterIO;
  type Params = CounterParams;
  type Inst = Counter;

  fn as_context(params: CounterParams, span: Option<MySpan>) -> Self {
    Self {
      module: cmtrs::gen_ir::ModuleMold::new(
        format!("Counter_{}", params.id_str()),
        span,
      ),
      params,
    }
  }

  fn module_mut(&mut self) -> &mut cmtrs::gen_ir::ModuleMold {
    &mut self.module
  }

  fn module(&self) -> &cmtrs::gen_ir::ModuleMold {
    &self.module
  }

  fn methods(&self) -> &'static [&'static str] {
    &["read", "set"]
  }

  fn params(&self) -> &Self::Params {
    &self.params
  }

  fn io(__cmt_gen: Rc<RefCell<Option<Self>>>) -> CounterIO {
    let set_val = input!("set_val".to_string(), {
      let t = __cmt_gen.borrow().as_ref().unwrap().params().t.clone();
      t
    });
    let count = output!("count".to_string(), {
      let t = __cmt_gen.borrow().as_ref().unwrap().params().t.clone();
      t
    });
    CounterIO { set_val, count }
  }

  // fn instantiate_subs(ctx: &Rc<RefCell<Option<Self>>>) {}

  fn into_inner(self) -> (gen_ir::ModuleMold, Self::Params) {
    (self.module, self.params)
  }
}

impl Instance for Counter {
  type Params = CounterParams;

  fn type_name() -> &'static str {
    &"Counter"
  }

  fn from_itfc(module: gen_ir::ModuleMold, params: Self::Params) -> Self {
    Self {
      module,
      ctx_view: Box::new(InstanceViewStruct::<CounterItfc>::new()),
      t: params.t,
    }
  }

  fn as_instance(&mut self, views: &mut InstanceViewIter) {
    self.ctx_view = views.next().unwrap();
  }

  fn as_nested_itfc_io(
    &self,
    name: String,
    views: &mut InstanceViewIter,
  ) -> Self {
    let instance = Self {
      module: gen_ir::ModuleMold::clone_with_new_name(&self.module, name),
      ctx_view: views.next().unwrap(),
      t: self.t.clone(),
    };
    // let set_val = module.add_input(format!("{}_set_val", name),
    // self.t.clone()); let count = module.add_output(format!("{}_count",
    // name), self.t.clone()); module.begin_nested_method(
    //   format!("{}_read", name),
    //   vec![],
    //   vec![count.value_id().unwrap()],
    // );
    // instance.read();
    // module.end_rule();
    // module.begin_nested_method(
    //   format!("{}_set", name),
    //   vec![set_val.value_id().unwrap()],
    //   vec![],
    // );
    // instance.set(set_val);
    // module.end_rule();
    instance
  }

  fn make_nested_methods(
    &self,
    name: String,
    __cmt_gen: ContextItfc<impl Itfc>,
  ) {
    let set_val = input!(format!("{}_set_val", name), self.t.clone());
    let count = output!(format!("{}_count", name), self.t.clone());
    named_method! {format!("{}_read", name); () -> (count) {self.read()}};
    named_method! {format!("{}_set", name); (set_val) {self.set(set_val)}};
  }

  fn module_mut(&mut self) -> &mut cmtrs::gen_ir::ModuleMold {
    &mut self.module
  }

  fn module(&self) -> &cmtrs::gen_ir::ModuleMold {
    &self.module
  }

  fn nested_names(&self) -> Vec<String> {
    vec![]
  }

  fn invoke(&self, name: &str, args: &[&cmtrs::Var], num_res: usize, span: Option<MySpan>) -> Vars {
    Vars(self.ctx_view.invoke(
      name.to_string(),
      args.iter().map(|x| x.ast()).collect(),
      num_res,
      span,
    ))
  }
}

impl Counter {
  pub fn read(&self) -> Var {
    self
      .ctx_view
      .invoke("read".to_string(), Vec::new(), 1, None)
      .into_iter()
      .next()
      .unwrap()
  }

  pub fn set(&self, a: impl CmtAST) {
    self.ctx_view.invoke("set".to_string(), vec![a.ast()], 0, None);
  }
}

impl ToIOMacro for Counter {
  type Output = Self;
  fn to_io_macro(self) -> Self::Output {
    self
  }
}
// *************************************

// Counter generater, contents should be generated by proc macro
//#[module]
fn make_counter(lim: usize, ty: &Type) -> Counter {
  let __cmt_gen = Rc::new(RefCell::new(Some(CounterItfc::as_context(
    CounterParams { t: ty.clone() },
    None,
  ))));

  let io = Itfc::io(__cmt_gen.clone());

  let reg_i = Itfc::instantiate(&__cmt_gen, "reg_i".to_string(), stl::reg(ty));

  // rule inc always {
  //   if i < lim { i = i + 1 } else { i = 0 }
  // }
  Itfc::begin_rule(
    &__cmt_gen,
    false,
    false,
    "inc".to_string(),
    cmtir::RuleSignature::Always {
      inputs: vec![],
      outputs: vec![],
    },
    cmtir::RuleTiming::SingleCycle,
    None,
    None,
    None,
  );

  Itfc::rule_guard(&__cmt_gen);

  Itfc::if_(&__cmt_gen, None);

  Itfc::statement(&__cmt_gen, reg_i.read().lt(literal(lim as i32, ty)));
  Itfc::sep(&__cmt_gen);

  Itfc::block(&__cmt_gen, None);
  let x = Itfc::add_var(&__cmt_gen, Some("x".to_string()));

  Itfc::assign(&__cmt_gen, &x, reg_i.read() + literal(1, ty), None);
  reg_i.write(&x);

  Itfc::statement(&__cmt_gen, Itfc::end(&__cmt_gen));
  Itfc::sep(&__cmt_gen);

  reg_i.write(literal(0, ty));

  Itfc::statement(&__cmt_gen, Itfc::end(&__cmt_gen));

  let _inc = Itfc::end_rule(&__cmt_gen);

  // rule set(set_val) {
  //    i = set_val
  // }
  Itfc::begin_rule(
    &__cmt_gen,
    false,
    false,
    "set".to_string(),
    cmtir::RuleSignature::Method {
      inputs: vec![io.set_val.value_id().unwrap()],
      outputs: vec![],
      side_effect: Some(true),
    },
    cmtir::RuleTiming::SingleCycle,
    None,
    None,
    None,
  );

  Itfc::rule_guard(&__cmt_gen);

  reg_i.write(io.set_val);

  let set = Itfc::end_rule(&__cmt_gen);

  // rule read() -> (count) {
  //    i.read()
  // }
  Itfc::begin_rule(
    &__cmt_gen,
    false,
    false,
    "read".to_string(),
    cmtir::RuleSignature::Method {
      inputs: vec![],
      outputs: vec![io.count.value_id().unwrap()],
      side_effect: Some(false),
    },
    cmtir::RuleTiming::SingleCycle,
    None,
    None,
    None,
  );

  Itfc::rule_guard(&__cmt_gen);

  Itfc::statement(&__cmt_gen, reg_i.read().ast());

  let read = Itfc::end_rule(&__cmt_gen);

  // Itfc::method_rel(&mut _cmt_gen, MethodRel::SA, &[set], &[read]);
  Itfc::method_rel(&__cmt_gen, MethodRel::C, &[set], &[set]);

  Itfc::dump(__cmt_gen)
}

fn main() {
  let ty = Type::UInt(4);
  let counter = make_counter(10, &ty);
  let circuit = counter.to_cmtir();
  println!("{}", circuit.ir_dump());
}