module type Droppable = sig
  type t
  val drop : t -> unit (* must not raise *)
end

module Droppable_in_channel : Droppable with type t = in_channel
  = struct
  type t = in_channel
  let drop = close_in_noerr
end

module type Affine = sig
  type t
  (* owning type *)
  type t_ref
  (* associated reference type *)
  val create : t_ref -> t
  (* Create a resource *)
  val borrow : t -> t_ref
  (* Borrow a resource. Raises Use_after_move if the resource has been
     moved or destroyed meanwhile. If the borrowed value is used after
     the resource has been destroyed, this is an error which cannot be
     detected at runtime. In a runtime that allocates memory with
     RAII, this can segfault, but here everything is GC'd so it will
     only violate user's invariants (we give an example later). *)
  val move : t -> t
  (* Move a resource. Using it on a non-live resource raises
     Use_after_free. Any copy not encapsulated in move can result in
     Double_free. *)
  val delete : t -> unit
  (* Compiler-generated destructor. Only to be called by RAII.scope. *)
end

module RAII : sig
  type 'a ptr
  module Make (M : Droppable) : (Affine with type t_ref = M.t)
  (* make a new resource type from a droppable type *)
  val handle : (unit -> 'a) -> 'a
  (* Ensures that destructors are called in order when an exception
     escapes. The bodies of try..withs must be wrapped in a call to
     handle. Failure to do so results in leaks when raising an
     exception. *)
  val scope : (module Affine with type t = 'a) -> 'a -> ('a -> 'b) -> 'b
  (* Simulate a scope for a bound owned variable. Any resource not
     encapsulated in such a scope leaks. *)
end
  = struct
  type 'a ptr = 'a cell ref
  and 'a cell = Live of 'a | Moved | Freed
  type destructor_closure = (unit -> unit) Stack.t
  let dcs : destructor_closure Stack.t = Stack.create ()
  module S = Stack
  let push_closure () =
    S.push (S.create ()) dcs
  let destroy_closure () =
    S.iter (fun f -> f ()) (S.pop dcs)
  let push_destructor (f : unit -> unit) =
    S.push f (S.top dcs)
  let pop_destructor () : unit -> unit =
    S.pop (S.top dcs)
  let handle (f : unit -> 'a) : 'a =
    push_closure ();
    match f () with
    | x -> destroy_closure (); x
    | exception e -> destroy_closure (); raise e
  let _ = push_closure (); at_exit destroy_closure
  module Make (M : Droppable) = struct
    type t = M.t ptr
    type t_ref =  M.t
    let create x = ref (Live x)
    let borrow o = match !o with
      | Live x -> x
      | Moved ->
         failwith "use after move"
         (* This can happen when borrowing after a move. Avoided using:
            a linearity checker. *)
      | Freed ->
         failwith "use after free"
         (* This can happen after copying a resource. Avoided by: the
            language, by always moving. This does not account for all
            use-after-free bugs: the hard ones, which we cannot detect at
            runtime, are obtained when the reference is created during the
            lifetime of the resource, but used after it is destroyed. The
            latter has to be avoided with a lifetime checker. *)
    let move o =
      let o' = create (borrow o) in o := Moved; o'
    let delete o = match !o with
      | Live x -> begin
          o := Freed;
          try M.drop x with
            _ -> ()
            (* An exception is raised while there could already stack
               unwinding for another raised exception. One can either
               call it undefined behaviour as in C++ and panic, or
               specify that all exceptions raise in destructors are
               silently ignored. *)
        end
      | Moved -> ()
      | Freed ->
         failwith "double free"
         (* This can happen after copying a resource. Avoided by: the
            language, by always moving. *)
  end
  let scope (type a) (module M : Affine with type t = a) x f =
    push_destructor (fun () -> M.delete x);
    let r = f x in
    pop_destructor () ();
    r
end

module File_in = RAII.Make (Droppable_in_channel)
type file_in = File_in.t
let scope_file_in f = RAII.scope (module File_in) f

let open_file name : file_in = File_in.create (open_in name)

let drop x =
  scope_file_in x @@ fun x ->
  ()

let drop2' x y =
  scope_file_in x @@ fun x ->
  scope_file_in y @@ fun y ->
  ()

let fancy_drop x =
  scope_file_in x @@ fun x ->
  try
    RAII.handle @@ fun () ->
    scope_file_in (File_in.move x) @@ fun y ->
    raise Exit
  with
  | Exit -> ()


let create_and_move name =
  scope_file_in (open_file name) @@ fun f ->
  drop (File_in.move f)

let twice1 name =
  scope_file_in (open_file name) @@ fun f ->
  (File_in.move f, File_in.move f)
(* failure "use after move" *)

let twice2 name =
  scope_file_in (open_file name) @@ fun f ->
  drop (File_in.move f);
  File_in.move f
(* failure "use after move" *)

let twice3 name =
  scope_file_in (open_file name) @@ fun f ->
  drop f; File_in.move f
(* failure "use after free". *)
let twice4 name =
  scope_file_in (open_file name) @@ fun f ->
  drop f; drop f
(* failure "double free". *)


let read_line name =
  scope_file_in (open_file name) @@ fun f ->
  print_endline (input_line (File_in.borrow f));
  (* if input_line raises an exception, f is closed then *)
  flush stdout
  (* f is closed then *)

let use_after_free name =
  scope_file_in (open_file name) @@ fun f ->
  let g = File_in.borrow f in
  drop (File_in.move f);
  input_line g (* Exception: Sys_error "Bad file descriptor". *)


module ATensor (P : Affine) (Q : Affine) : Affine
       with type t = P.t * Q.t
       with type t_ref = P.t_ref * Q.t_ref
  = struct
  type t = P.t * Q.t
  type t_ref = P.t_ref * Q.t_ref
  let create (x,y) = (P.create x, Q.create y)
  let borrow (x,y) = (P.borrow x, Q.borrow y)
  let move (x,y) = (P.move x, Q.move y)
  (* What about all the deep copies? In the actual proposed runtime model,
     moves and borrows only involve copying or moving the pointer. *)
  let delete (x,y) = P.delete x; Q.delete y
end

let open2 name1 name2 =
  let module File2 = ATensor (File_in) (File_in) in
  scope_file_in (open_file name2) @@ fun y ->
  scope_file_in (open_file name1) @@ fun x ->
  File2.move (x,y)


module Affine_of_GCd (M : sig type t end) : Affine
       with type t = M.t
       with type t_ref = M.t
  = struct
  type t = M.t
  type t_ref = M.t
  let create x = x
  let borrow x = x
  let move x = x
  let delete x = ()
end

module Affine_of_String =
  Affine_of_GCd (struct type t = string end)
module String_and_file =
  ATensor (Affine_of_GCd (struct type t = string end)) (File_in)
let fst (z : string * file_in) : string =
  RAII.scope (module String_and_file) z @@ fun (x,y) ->
  x

let snd (z : string * file_in) : file_in =
  RAII.scope (module String_and_file) z @@ fun (x,y) ->
  File_in.move y


type ('a,'b) sum = Left of 'a | Right of 'b

module ASum (P : Affine) (Q : Affine)
       : Affine
       with type t = (P.t, Q.t) sum
       with type t_ref = (P.t_ref, Q.t_ref) sum
  = struct
  type t = (P.t, Q.t) sum
  type t_ref = (P.t_ref, Q.t_ref) sum
  let borrow = function
    | Left x -> Left (P.borrow x)
    | Right y -> Right (Q.borrow y)
  let create = function
    | Left x -> Left (P.create x)
    | Right y -> Right (Q.create y)
  let move = function
    | Left x -> Left (P.move x)
    | Right y -> Right (Q.move y)
  let delete = function
    | Left x -> P.delete x
    | Right y -> Q.delete y
end

module From_map (M : sig
             type 'a t
             val map : ('a -> 'b) -> 'a t -> 'b t
           end) (P : Affine) : Affine
       with type t = P.t M.t
       with type t_ref = P.t_ref M.t
  = struct
  type t = P.t M.t
  type t_ref = P.t_ref M.t
  let create x = M.map P.create x
  let borrow x = M.map P.borrow x
  let move x = M.map P.move x
  let delete x = let _ = M.map P.delete x in ()
end

module AOption = From_map (struct
                     type 'a t = 'a option
                     let map f = function
                       | Some x -> Some (f x)
                       | None -> None
                   end)

module AList = From_map (struct
                   type 'a t = 'a list
                   let map = List.map (* destroys in LIFO order *)
                 end)



module type Mutex_sig = sig
  type t
  val create : unit -> t
  val lock : t -> unit
  val try_lock : t -> bool
  val unlock : t -> unit
end

module RAII_Mutex(Mutex : Mutex_sig) : sig
  type t
  type lock
  module Lock : Affine with type t = lock
  val create : unit -> t
  val lock : t -> lock
  val try_lock : t -> lock option
end = struct
  type t = Mutex.t
  module Lock = RAII.Make(struct
                    type t = Mutex.t
                    let drop = Mutex.unlock
                  end)
  type lock = Lock.t
  let create = Mutex.create
  let lock m = Mutex.lock m; Lock.create m
  let try_lock m = if Mutex.try_lock m then
                     Some(Lock.create m) (* scope is unnecessary
                                            because the context is
                                            pure *)
                   else None
end

module Locking (Mutex : Mutex_sig) : sig
  type t
  module T : Affine with type t = t
  val try_lock : RAII_Mutex(Mutex).t list -> T.t option
end
  = struct
  module M = RAII_Mutex (Mutex)
  module T = AList (M.Lock)
  module O = AOption (M.Lock)
  type t = M.lock list
  let scope x = RAII.scope (module T) x
  let try_lock ms =
    let try_with_exn (m : M.t) : M.lock =
      RAII.scope (module O) (M.try_lock m) @@ function
       | Some l -> M.Lock.move l
       | None -> raise Exit
    in
    try RAII.handle (fun () -> Some (List.rev_map try_with_exn ms))
    (* Wrong: assumes a ressource-aware List.rev_map, that releases
       partial lockings in case of failure. It is given further below. *)
    with Exit -> None
end

module CapArray : sig
  type ('a,'b) t
  type 'b cap
  type 'b cap_ref (* One cannot do anything with a &cap, but we
                     need it for defining &cap_array *)
  module Cap (B : sig type b end) : Affine
         with type t = B.b cap
         with type t_ref = B.b cap_ref
  type _ cap_array = Cap_array : ('a,'b) t * 'b cap -> 'a cap_array
  type _ cap_array_ref = Cap_array_ref : ('a,'b) t * 'b cap_ref -> 'a cap_array_ref
  module Cap_Array (A : sig type a end) : Affine
         with type t = A.a cap_array
         with type t_ref = A.a cap_array_ref
  val make : int -> 'a -> 'a cap_array
  val set : ('a,'b) t -> int -> 'a -> 'b cap -> 'b cap
  val get : ('a,'b) t -> int -> 'b cap -> 'a * 'b cap
  val dirty_get : ('a,'b) t -> int -> 'a
end = struct
  type ('a,'b) t = 'a array
  type _ cap = unit
  type 'a cap_ref = 'a cap
  module Cap (B : sig type b end) = struct
    type t = unit
    type t_ref = unit
    let create () = ()
    let borrow () = ()
    let move () = ()
    let delete () = () (* a capability is affine and trivially destructible *)
  end
  type _ cap_array = Cap_array : ('a,'b) t * 'b cap -> 'a cap_array
  type _ cap_array_ref = Cap_array_ref : ('a,'b) t * 'b cap_ref -> 'a cap_array_ref
  let f : 'a cap -> (module Affine with type t = 'a cap) = fun (type a) x ->
    let module M = Cap (struct type b = a end) in
    (module M : Affine with type t = a cap)
  module Cap_Array (A : sig type a end) = struct
    type t = A.a cap_array
    type t_ref = A.a cap_array_ref
    let create = function Cap_array_ref (x,c) -> Cap_array (x,c)
    let borrow = function Cap_array (x,c) -> Cap_array_ref (x,c)
    let move x = x
    let delete x = ()
  end
  let make n (type a') (x : a') =
    let module M = Cap_Array (struct type a = a' end) in
    M.move (Cap_array (Array.make n x,()))
  let set ar n x () = Array.set ar n x
  let get ar n () = (Array.get ar n, ())
  let dirty_get ar n = Array.get ar n
end


let drop (type a) (module A : Affine with type t = a) (x : a) =
  RAII.scope (module A) x @@ fun x ->
  ()

let merge (type a) (module A : Affine with type t = a) =
  let module AL = AList (A) in
  let module AL2 = ATensor (AL) (AL) in
  let rec iter x =
    RAII.scope (module AL2) x @@ function
     | list, [] | [], list -> AL.move list
     | h1::t1, h2::t2 ->
        if A.borrow h1 <= A.borrow h2 then
          A.move h1 :: iter (AL.move t1, AL.move (h2::t2))
        else
          A.move h2 :: iter (AL.move (h1::t1), AL.move t2)
  in
  iter


module Mutable (P : Affine) : sig
  module T : Affine
  val ref : P.t -> T.t
  val (:=) : T.t_ref -> P.t -> unit
  val (!) : T.t -> P.t
  val (!&) : T.t_ref -> P.t_ref
  val swap : T.t_ref -> T.t_ref -> unit
end = struct
  type t' = { mutable contents: P.t }
  module T = struct
    type t = t'
    type t_ref = t
    let create r = r
    let borrow r = r
    let move r = { contents = P.move r.contents }
    let delete r = P.delete r.contents
  end
  let ref x = { contents = P.move x }
  let (!) r = P.move r.contents
  let (!&) r = P.borrow r.contents
  let (:=) r x =
    RAII.scope (module P) (!r) @@ fun _ ->
    r.contents <- P.move x
  let swap r r' =
    let x = r.contents in
    r.contents <- r'.contents;
    r'.contents <- x
end


module AFun : sig
  type ('a,'b) fn
  type ('a,'b) fn_ref
  type ('a,'b) fn_once
  module Fn (M : sig type a type b end) : Affine
         with type t = (M.a, M.b) fn
         with type t_ref = (M.a, M.b) fn_ref
  module FnOnce (M : sig type a type b end) : Affine
         with type t = (M.a, M.b) fn_once
  val make_once : ('a -> 'b) -> ('a,'b) fn_once
  val move_into_once : (module Affine with type t = 'c) ->
                       'c -> ('c -> ('a,'b) fn_once) -> ('a,'b) fn_once
  val move_into : (module Affine with type t = 'c
                                  and type t_ref = 'd) ->
                  'c -> ('d -> ('a,'b) fn) -> ('a,'b) fn
  val app_ref : ('a,'b) fn_ref -> 'a -> 'b
  val app_once : ('a,'b) fn_once -> 'a -> 'b
end = struct
  module type Closure = sig
    module A : Affine
    val x : A.t
  end
  let closure_move (module M : Closure) = (module struct
                                             module A = M.A
                                             let x = A.move M.x
                                           end : Closure)
  let closure_delete (module M : Closure) = M.(A.delete x)
  type ('a,'b) fn = ('a -> 'b) * (module Closure) list
  type ('a,'b) fn_ref = 'a -> 'b
  type ('a,'b) fn_once = ('a,'b) fn
  module Fn (M : sig type a type b end) : Affine
         with type t = (M.a, M.b) fn
         with type t_ref = (M.a, M.b) fn_ref
    = struct
    type t = (M.a, M.b) fn
    type t_ref = M.a -> M.b
    let create f = (f,[])
    let borrow (f,l) = f
    let move (f,l) = (f, List.map closure_move l)
    let delete (f,l) = List.iter closure_delete l
  end
  module FnOnce (M : sig type a type b end) : Affine
         with type t = (M.a, M.b) fn_once
    = Fn (M)
  let make_once f = (f,[])
  let move_into_once (type c) (module A : Affine with type t = c)
        res closure =
    let x = A.move res in
    let (f,l) = closure x in
    let x_mod = (module struct
                   module A = A
                   let x = x
                 end : Closure)
    in
    (f,x_mod::l)
  let move_into (type c) (type d)
        (module A : Affine with type t = c and type t_ref = d)
        res closure =
    let x = A.move res in
    let (f,l) = closure (A.borrow x) in
    let x_mod = (module struct
                   module A = A
                   let x = x
                 end : Closure)
    in
    (f,x_mod::l)
  let app_ref = (@@)
  let app_once (type a') (type b') fn x =
    let module F = Fn (struct type a = a' type b = b' end) in
    RAII.scope (module F) fn @@ fun (f,_) ->
    f x
end

let drop2 x =
  scope_file_in x @@ fun x ->
  AFun.move_into_once (module File_in) x (fun x ->
      AFun.make_once (fun y ->
          (scope_file_in y @@ fun y -> ())))
(* : File_in.t -> (File_in.t, unit) AFun.fn_once *)


let rev_map (type a) (module A : Affine with type t = a)
      (type b) (module B : Affine with type t = b)
      (f : a -> b) l =
  let module AL = AList(A) in
  let module BL = AList(B) in
  let exception Tail_rec of BL.t * AL.t in
  let rmap_f accu x =
    RAII.scope (module BL) accu @@ fun accu ->
    RAII.scope (module AL) x @@ function
     | [] -> BL.move(accu)
     | a::l -> raise (Tail_rec (f (A.move a) :: BL.move(accu), AL.move l))
  in
  let rec tail_rec x y =
    try
      RAII.handle (fun () -> rmap_f x y)
    with
      Tail_rec (x',y') -> (tail_rec [@tailcall]) x' y'
  in
  tail_rec [] l


let test_rev_map () =
  let x = open_file "/tmp/dummy1" in
  let y = open_file "/tmp/dummy2" in
  let z = open_file "/tmp/dummy3" in
  let z' = File_in.borrow z in
  let rev_files = rev_map (module File_in) (module File_in)
                    (fun x -> File_in.move x)
  in
  let module L = AList(File_in) in
  let scope = RAII.scope (module L) in
  let _ =
      scope [File_in.move x; File_in.move y; File_in.move z] @@ fun l ->
      scope (rev_files (L.move l)) @@ fun l' ->
      scope (rev_files (L.move l')) @@ fun l'' ->
      (* val l'' : File_in.t list = [<abstr>; <abstr>; <abstr>] *)
      rev_map (module File_in) (module Affine_of_GCd(struct type t = unit end))
        (fun x -> scope_file_in x @@
         fun x -> (x |> File_in.borrow |> input_line |> print_endline))
        (L.move l'')
  in
  (* check that files are closed *)
  print_endline (input_line z') (* Exception: Sys_error "Bad file descriptor". *)


let swap x y = let z = !x in x := !y; y := z

(* Rc_Cache for GCed types. Completely useless! *)
module Rc_Cache (Ord : Map.OrderedType) : sig
  type key = Ord.t
  type 'a t
  type 'a cached
  val create : unit -> 'a t
  val get_or_add : key -> (unit -> 'a) -> 'a t -> 'a cached
  val get : key -> 'a t -> ('a cached) option
  val deref : 'a cached -> 'a
  val copy : 'a cached -> 'a cached
end = struct
  type key = Ord.t

  module M = Map.Make (Ord)

  type 'a t = ('a * int ref) M.t ref

  let create () = ref M.empty

  let take cache =
    let cache' = create() in
    swap cache cache';
    !cache'

  let decr k cache =
    match M.find k !cache with
    | (_, r) -> begin
        decr r;
        if !r = 0 then
          cache := take cache |> M.remove k
      end
    | exception Not_found -> assert false

  type 'a cached = 'a * key * 'a t

  let deref (x, _, _) = x

  let get k cache =
    match M.find k !cache with
    | (x, r) -> begin
        incr r;
        Some (x, k, cache)
      end
    | exception Not_found -> None

  let get_unsafe k cache =
    match get k cache with
    | Some cached -> cached
    | None -> assert false

  let get_or_add k f cache =
    match get k cache with
    | Some cached -> cached
    | None -> begin
        cache := take cache |> M.add k (f(), ref 0);
        get_unsafe k cache
      end

  let copy (_,k,cache) =
    get_unsafe k cache
end