(* move / affinity 검사.

   보안 정리 (ii) — safe code에서 capability는 복제·위조되지 않는다 — 를 코드로
   닫는 단계다. 검사는 전부 함수 로컬 데이터플로우다. 전역 분석이 없다.

   affinity의 뿌리는 capability다. capability를 필드로 가진 타입은 전이적으로
   affine이고(철학 3), 이 전이가 없으면 wrapper 하나를 복사해 capability가
   사실상 복제된다.

   외부 타입은 affine임을 증명할 수 없으므로 copyable로 본다 — 모르는 것을
   위반이라고 말하지 않는다. 그래서 이 검사를 시험하려면 자원 타입을 모듈
   안에서 정의해야 한다 (samples/05). *)

open Ast

type error = { pos : Token.pos; msg : string }

(* 값이 무엇인가: 소유한 affine 값인가, 빌린 값인가. 둘은 직교한다. *)
type vinfo = { v_affine : bool; v_use : bool }

let v_copy = { v_affine = false; v_use = false }

type binding = {
  b_id : int;
  b_name : string;
  b_affine : bool;
  b_use : bool;
  b_mut : bool;
  b_depth : int;
}

type fninfo = { f_params : param list; f_ret : Ast.ty option }

type state = {
  aff : (string, bool) Hashtbl.t;
  fns : (string, fninfo) Hashtbl.t;
  meths : (string, (string * fninfo) list) Hashtbl.t;
  mutable scopes : binding list list;
  moved : (int, Token.pos) Hashtbl.t;
  mutable next_id : int;
  mutable depth : int;
  (* 클로저 프레임: (프레임 깊이, 잡아온 바깥 바인딩). 중첩 클로저를 위해 스택 *)
  mutable frames : (int * binding list ref) list;
  mutable errors : error list;
}

let err st pos msg = st.errors <- { pos; msg } :: st.errors

(* ------------------------------------------------------------------ *)
(* affinity 유도                                                       *)
(* ------------------------------------------------------------------ *)

let builtin_containers = [ "List"; "Option"; "Result" ]

let rec ty_affine st (t : Ast.ty) =
  match t with
  | T_fn { affine; _ } -> affine
  | T_named { modl; name; args; _ } ->
      let name = match modl with Some a -> a ^ "." ^ name | None -> name in
      let self =
        match Hashtbl.find_opt st.aff name with Some b -> b | None -> false
      in
      let arg_affine =
        List.exists
          (function TA_ty t -> ty_affine st t | TA_eff _ -> false)
          args
      in
      self || (List.mem name builtin_containers && arg_affine) || arg_affine

(* capability가 뿌리다. struct/enum은 필드에서 전이된다. 상호 재귀 타입을 위해
   변화가 없을 때까지 돈다. *)
let derive_affinity st (items : item list) =
  List.iter
    (fun it ->
      match it with
      | I_capability { name; _ } -> Hashtbl.replace st.aff name true
      | I_struct { name; _ } -> Hashtbl.replace st.aff name false
      | I_enum { name; _ } -> Hashtbl.replace st.aff name false
      | _ -> ())
    items;
  let changed = ref true in
  while !changed do
    changed := false;
    List.iter
      (fun it ->
        let update name affine =
          if affine && Hashtbl.find_opt st.aff name <> Some true then (
            Hashtbl.replace st.aff name true;
            changed := true)
        in
        match it with
        | I_struct { name; fields; _ } ->
            update name (List.exists (fun f -> ty_affine st f.f_ty) fields)
        | I_enum { name; variants; _ } ->
            update name
              (List.exists
                 (fun v -> List.exists (ty_affine st) v.v_args)
                 variants)
        | _ -> ())
      items
  done;
  (* copyable 선언과 affine 필드는 공존할 수 없다 *)
  List.iter
    (fun it ->
      match it with
      | I_struct { copyable = true; name; fields; pos; _ } ->
          List.iter
            (fun f ->
              if ty_affine st f.f_ty then
                err st f.f_pos
                  (Printf.sprintf "%s은(는) copyable로 선언되었지만 %s 필드가 affine입니다"
                     name f.f_name))
            fields;
          ignore pos
      | _ -> ())
    items

(* ------------------------------------------------------------------ *)
(* 스코프                                                              *)
(* ------------------------------------------------------------------ *)

let push st = st.scopes <- [] :: st.scopes
let pop st = match st.scopes with _ :: r -> st.scopes <- r | [] -> ()

let add st name ~affine ~use ~mut_ =
  st.next_id <- st.next_id + 1;
  let b =
    {
      b_id = st.next_id;
      b_name = name;
      b_affine = affine;
      b_use = use;
      b_mut = mut_;
      b_depth = st.depth;
    }
  in
  (match st.scopes with
  | s :: r -> st.scopes <- (b :: s) :: r
  | [] -> st.scopes <- [ [ b ] ]);
  b

let find st name =
  let rec go = function
    | [] -> None
    | s :: r -> (
        match List.find_opt (fun b -> b.b_name = name) s with
        | Some b -> Some b
        | None -> go r)
  in
  go st.scopes

(* 클로저 안에서 바깥 바인딩을 건드리면 capture다. 프레임마다 기록한다. *)
let note_capture st b =
  List.iter
    (fun (fdepth, acc) ->
      if b.b_depth < fdepth && not (List.exists (fun x -> x.b_id = b.b_id) !acc)
      then acc := b :: !acc)
    st.frames

(* ------------------------------------------------------------------ *)
(* 분기 병합 — 보수적 합집합                                            *)
(* ------------------------------------------------------------------ *)

let snapshot st = Hashtbl.copy st.moved

let restore st snap =
  Hashtbl.reset st.moved;
  Hashtbl.iter (fun k v -> Hashtbl.replace st.moved k v) snap

let merge st snaps =
  (* 한 분기에서라도 moved면 병합 지점 이후 moved *)
  Hashtbl.reset st.moved;
  List.iter
    (fun snap ->
      Hashtbl.iter
        (fun k v ->
          if not (Hashtbl.mem st.moved k) then Hashtbl.replace st.moved k v)
        snap)
    snaps

(* ------------------------------------------------------------------ *)
(* 식                                                                  *)
(* ------------------------------------------------------------------ *)

type ctx = Borrow | Move of string (* 어디로 옮겨가는지 — 진단에 쓴다 *)

let rec walk st (ctx : ctx) (e : expr) : vinfo =
  match e with
  | E_lit _ -> v_copy
  | E_ident (n, pos) -> (
      match find st n with
      | None -> v_copy
      | Some b ->
          note_capture st b;
          (match Hashtbl.find_opt st.moved b.b_id with
          | Some mp ->
              err st pos
                (Printf.sprintf "%s은(는) 이미 move되었습니다 (%d:%d에서 소비)" n
                   mp.Token.line mp.Token.col)
          | None -> ());
          (match ctx with
          | Borrow -> ()
          | Move where ->
              if b.b_use then
                err st pos
                  (Printf.sprintf "%s은(는) 빌린 값이라 %s 없습니다 (own으로 받아야 합니다)" n
                     where)
              else if b.b_affine then Hashtbl.replace st.moved b.b_id pos);
          { v_affine = b.b_affine; v_use = b.b_use })
  | E_list (xs, _) ->
      let infos = List.map (walk st (Move "컨테이너에 넣을 수")) xs in
      { v_affine = List.exists (fun i -> i.v_affine) infos; v_use = false }
  | E_struct { name; fields; _ } ->
      List.iter (fun (_, e) -> ignore (walk st (Move "struct에 저장할 수") e)) fields;
      {
        v_affine =
          (match Hashtbl.find_opt st.aff name with
          | Some b -> b
          | None -> false);
        v_use = false;
      }
  | E_closure c -> (
      let r = walk_closure st c in
      (* use의 전염: 빌린 값을 capture한 클로저는 그 자체가 빌린 값이라
         소유를 가져가는 자리로 갈 수 없다. 별도의 nonescaping 개념 없이
         use 규칙 하나로 닫힌다. *)
      match ctx with
      | Move where when r.v_use ->
          err st c.cl_pos
            (Printf.sprintf "빌린 값을 capture한 클로저는 %s 없습니다 (use 값은 탈출하지 못합니다)"
               where);
          r
      | _ -> r)
  | E_if { cond; then_; else_; _ } ->
      ignore (walk st Borrow cond);
      let before = snapshot st in
      let t1 = walk_block st ctx then_ in
      let s1 = snapshot st in
      restore st before;
      let t2 = match else_ with None -> v_copy | Some e -> walk st ctx e in
      let s2 = snapshot st in
      merge st [ s1; s2 ];
      { v_affine = t1.v_affine || t2.v_affine; v_use = t1.v_use || t2.v_use }
  | E_match { scrutinee; arms; _ } ->
      let sinfo = walk st Borrow scrutinee in
      let before = snapshot st in
      let results =
        List.map
          (fun a ->
            restore st before;
            push st;
            bind_pattern st sinfo a.arm_pat;
            let r = walk st ctx a.arm_body in
            pop st;
            (r, snapshot st))
          arms
      in
      if results <> [] then merge st (List.map snd results);
      {
        v_affine = List.exists (fun (r, _) -> r.v_affine) results;
        v_use = List.exists (fun (r, _) -> r.v_use) results;
      }
  | E_scope { name; parent; body; _ } ->
      ignore (walk st Borrow (E_ident (parent, Token.{ line = 0; col = 0 })));
      push st;
      (* 자식 TaskScope는 second-class다 — 블록 밖으로 나갈 수 없다 *)
      ignore (add st name ~affine:true ~use:true ~mut_:false);
      let r = walk_block st ctx body in
      pop st;
      r
  | E_block b ->
      push st;
      let r = walk_block st ctx b in
      pop st;
      r
  | E_call { callee; args; pos } -> walk_call st callee args pos
  | E_field { obj; _ } ->
      (* v0에 부분 move는 없다. 필드 접근은 빌림이고 결과도 빌린 값이다. *)
      let o = walk st Borrow obj in
      if o.v_affine || o.v_use then { v_affine = true; v_use = true }
      else v_copy
  | E_inst { callee; _ } -> walk st Borrow callee
  | E_try { inner; _ } -> walk st ctx inner
  | E_unary { operand; _ } -> walk st Borrow operand
  | E_binary { lhs; rhs; _ } ->
      ignore (walk st Borrow lhs);
      ignore (walk st Borrow rhs);
      v_copy

and bind_pattern st (info : vinfo) p =
  match p with
  | P_wild _ | P_lit _ -> ()
  | P_bind (n, _) ->
      ignore (add st n ~affine:info.v_affine ~use:info.v_use ~mut_:false)
  | P_ctor { args; _ } -> List.iter (bind_pattern st info) args

and walk_block st ctx (b : block) : vinfo =
  let rec go = function
    | [] -> v_copy
    | [ S_expr e ] -> walk st ctx e (* 꼬리 식은 블록의 값이다 *)
    | s :: rest ->
        walk_stmt st s;
        go rest
  in
  go b.stmts

and walk_stmt st = function
  | S_let { mut_; pat; value; _ } ->
      let info = walk st (Move "다른 이름에 묶을 수") value in
      let rec bind p =
        match p with
        | P_bind (n, _) ->
            ignore (add st n ~affine:info.v_affine ~use:info.v_use ~mut_)
        | P_ctor { args; _ } -> List.iter bind args
        | _ -> ()
      in
      bind pat
  | S_return { value; _ } -> (
      match value with
      | None -> ()
      | Some e -> ignore (walk st (Move "반환할 수") e))
  | S_assign { place; value; _ } ->
      ignore (walk st (Move "대입할 수") value);
      ignore (walk st Borrow place)
  | S_expr e -> ignore (walk st Borrow e)

(* 클로저: 무엇을 잡아왔는지가 클로저 자신의 성질을 정한다 (전이 규칙) *)
and walk_closure st (c : closure) : vinfo =
  st.depth <- st.depth + 1;
  let acc = ref [] in
  st.frames <- (st.depth, acc) :: st.frames;
  push st;
  List.iter
    (fun (n, ann) ->
      let affine = match ann with Some t -> ty_affine st t | None -> false in
      ignore (add st n ~affine ~use:false ~mut_:false))
    c.cl_params;
  ignore (walk_block st (Move "반환할 수") c.cl_body);
  pop st;
  st.frames <- List.tl st.frames;
  st.depth <- st.depth - 1;
  let captured = !acc in
  let affine = ref false and use_ = ref false in
  List.iter
    (fun b ->
      if b.b_mut then
        err st c.cl_pos
          (Printf.sprintf
             "클로저는 mut 바인딩 %s을(를) capture할 수 없습니다 (v0에 참조가 없으므로 별칭도 복사도 만들지 \
              않는다)"
             b.b_name);
      if b.b_use then use_ := true;
      if b.b_affine && not b.b_use then (
        (* by-move capture: 바깥에서는 여기서 소비된다 *)
        affine := true;
        if not (Hashtbl.mem st.moved b.b_id) then
          Hashtbl.replace st.moved b.b_id c.cl_pos))
    captured;
  { v_affine = !affine; v_use = !use_ }

and walk_call st callee args pos =
  let info = callee_info st callee in
  ignore (walk st Borrow callee);
  let params = match info with Some f -> f.f_params | None -> [] in
  List.iteri
    (fun i a ->
      let p = List.nth_opt params i in
      let own = match p with Some p -> p.p_own | None -> false in
      let pty = Option.map (fun p -> p.p_ty) p in
      let ctx = if own then Move "다른 함수에 넘길 수" else Borrow in
      let got = walk st ctx a in
      (* callable affinity: affine 클로저를 fn 자리에 넘길 수 없다 *)
      match pty with
      | Some (T_fn { affine = false; _ }) when got.v_affine && not got.v_use ->
          err st pos
            "affine 값을 capture한 클로저는 fn 자리에 넘길 수 없습니다 (affine fn이어야 합니다)"
      | _ -> ())
    args;
  match info with
  | Some { f_ret = Some t; _ } -> { v_affine = ty_affine st t; v_use = false }
  | _ -> v_copy

and callee_info st callee =
  match callee with
  | E_ident (n, _) when find st n = None -> Hashtbl.find_opt st.fns n
  | E_inst { callee = E_ident (n, _); _ } when find st n = None ->
      Hashtbl.find_opt st.fns n
  | E_field { obj = E_ident (o, _); name; _ } -> (
      match find st o with
      | Some _ -> (
          (* 값의 메서드: 타입을 모르면 넘어간다. capability 메서드는 아래에서 *)
          match
            Hashtbl.fold
              (fun _ methods acc ->
                match acc with
                | Some _ -> acc
                | None -> List.assoc_opt name methods)
              st.meths None
          with
          | Some f -> Some f
          | None -> None)
      (* 모듈 별칭을 통한 호출: 가져온 함수는 "Alias.f" 키로 들어와 있다 *)
      | None -> Hashtbl.find_opt st.fns (o ^ "." ^ name))
  | _ -> None

(* ------------------------------------------------------------------ *)
(* 선언                                                                *)
(* ------------------------------------------------------------------ *)

let check_fn st (d : fn_decl) =
  match d.fn_body with
  | None -> ()
  | Some body ->
      st.scopes <- [];
      Hashtbl.reset st.moved;
      st.depth <- 0;
      st.frames <- [];
      push st;
      List.iter
        (fun p ->
          let affine = ty_affine st p.p_ty in
          (* 무표기 = 빌림. own만이 소유 이전이다. *)
          let use_ = affine && not p.p_own in
          ignore (add st p.p_name ~affine ~use:use_ ~mut_:p.p_mut))
        d.fn_params;
      let r = walk_block st (Move "반환할 수") body in
      (match d.fn_ret with
      | Some (T_fn { affine = false; _ }) when r.v_affine && not r.v_use ->
          err st d.fn_pos
            (Printf.sprintf
               "%s이(가) affine 값을 capture한 클로저를 fn 타입으로 반환합니다 (affine fn이어야 합니다)"
               d.fn_name)
      | _ -> ());
      pop st

let check ?(imports : item list = []) (m : modul) : error list =
  let st =
    {
      aff = Hashtbl.create 16;
      fns = Hashtbl.create 16;
      meths = Hashtbl.create 16;
      scopes = [];
      moved = Hashtbl.create 16;
      next_id = 0;
      depth = 0;
      frames = [];
      errors = [];
    }
  in
  derive_affinity st (imports @ m.items);
  List.iter
    (fun it ->
      match it with
      | I_fn { decl; _ } ->
          Hashtbl.replace st.fns decl.fn_name
            { f_params = decl.fn_params; f_ret = decl.fn_ret }
      | I_capability { name; methods; _ } ->
          Hashtbl.replace st.meths name
            (List.map
               (fun d ->
                 (d.fn_name, { f_params = d.fn_params; f_ret = d.fn_ret }))
               methods)
      | _ -> ())
    (imports @ m.items);
  List.iter
    (fun it -> match it with I_fn { decl; _ } -> check_fn st decl | _ -> ())
    m.items;
  List.sort
    (fun a b ->
      compare
        (a.pos.Token.line, a.pos.Token.col)
        (b.pos.Token.line, b.pos.Token.col))
    (List.rev st.errors)