Arrow.ml

module StringSet = Set.Make(String)

module type ACTION = sig
  type 'a t

  val execute : 'a t -> 'a
end

module type CACHE = sig
  type key = string list

  val store : key -> 'value -> unit
  val get: key -> 'value option
end

module Build (Action : ACTION)(Cache : CACHE) : sig

  type 'a t
  type ('a, 'b) fbuild

  val return : 'a -> 'a t
  val execute : 'a Action.t -> 'a t

  val pair : 'a t -> 'b t -> ('a * 'b) t
  val many : 'a t list -> 'a list t

  val bind : f:('a -> 'b t) -> 'a t -> 'b t

  val define : string -> 'a t -> 'a t

  val definef : ?cache:('a -> string) -> ('a -> 'b t) -> ('a, 'b) fbuild t
  val selectf : ('a, 'b) fbuild t -> 'a -> 'b t

  val run : 'a t -> 'a

  module Let_syntax : sig
    val bind : f:('a -> 'b t) -> 'a t -> 'b t
  end

end = struct

  (** Build *)
  type _ t =

    (** Primitives / atomic builds. *)

    (* Dummy build, does nothing, just returns what it has. *)
    | Return : 'a -> 'a t

    (* Execute action and returns its result. *)
    | Execute : 'a Action.t -> 'a t

    (** Parallel composition (applicative). *)

    (* Build which joins two other builds. *)
    | Pair : ('a t * 'b t) -> ('a * 'b) t
    (* Build which joins a list of builds. *)
    | Many : 'a t list -> 'a list t

    (** Sequential composition (monadic). *)

    (** Build which depends on another build's result. *)
    | Bind : {f : 'a -> 'b t; build : 'a t;} -> 'b t

    (** Named builds and build families. *)

    (** Build with an id. *)
    | Define : string * 'a t -> 'a t

    (** A family of builds indexed by id. *)
    | DefineF : ('a -> string) option * ('a -> 'b t) -> ('a, 'b) fbuild t
    (** Select build from a family of builds by passing an id. *)
    | SelectF : ('a, 'b) fbuild t * 'a -> 'b t

  and ('a, 'b) fbuild = FBuild of {
    cache : ('a -> string) option;
    f : 'a -> 'b t;
  }

  let return v = Return v
  let execute action = Execute action

  let pair a b = Pair (a, b)
  let many builds = Many builds

  let bind ~f build = Bind {f; build}

  let define name build = Define (name, build)

  let definef ?cache f = DefineF (cache, f)
  let selectf f id = SelectF (f, id)


  let run : type a. a t -> a = (fun build ->
      let rec run' : type a b. Cache.key -> a t -> a = fun path -> function
        | Return arg -> arg
        | Execute action -> Action.execute action
        | Pair (build1, build2) -> (run' path build1, run' path build2)
        | Many builds -> List.map (run' path) builds
        | Bind {f; build} -> run' path build |> f |> run' path
        | Define (id, build) ->
          begin
            match build with
            | Define _ -> failwith "Cannot nest define"
            | _ -> cached (path @ [id]) build
          end
        | DefineF (cache, f) ->
          FBuild {cache; f;}
        | SelectF (f, id) -> begin
            match run' path f with
            | FBuild {cache; f;} ->
                match cache with
                | None -> run' path (f id)
                | Some cache -> cached (path @ [cache id]) (f id)
          end
      and cached : type a b. Cache.key -> a t -> a = fun path build ->
        match Cache.get path with
        | Some value ->
          Printf.printf "HIT : %s\n" (String.concat "::" path);
          value
        | None ->
          Printf.printf "MISS: %s\n" (String.concat "::" path);
          let value = run' path build in
          Cache.store path value;
          value
      in
      run' [] build
  )

  module Let_syntax = struct
    let bind = bind
  end
end

module Fastpack = struct

  type path = string
  type jsdep = string
  type jssource = string
  type jsmodule = {id: string; source: jssource; path: path}
  type jsbundle = string

  module Cache : CACHE = struct
    type key = string list

    let cache = Hashtbl.create 100
    let get key = Obj.magic (Hashtbl.find_opt cache key)
    let store key value = Hashtbl.replace cache key (Obj.magic value)
  end

  module Action = struct

    (* Fastpack specific semantic actions. *)
    type _ t =
      | Read: path -> jssource t
      | Resolve: {
          context: jsmodule option;
          jsdep: jsdep
        } -> path t
      | Compile: path *jssource -> (jsmodule * jsdep list) t
      | Pack: jsmodule list -> jsbundle t

    (* Read and parse JS sources. *)
    let read : path -> jssource t = fun path -> Read path

    (* Resolves js dependency into a path. *)
    let resolve : ?context:jsmodule -> jsdep -> path t = fun ?context jsdep ->
      Resolve {context; jsdep}

    (* Compiels js sources and returns js module and a list of paths to deps *)
    let compile : path -> jssource -> (jsmodule * string list) t = fun path jssource ->
      Compile (path, jssource)

    (* Packs a list of js modules into a js bundle *)
    let pack : jsmodule list -> jsbundle t = fun modules -> Pack modules

    (* Total fake *)
    let execute : type res. res t -> res = function
      | Read path -> "CONTENTS OF: " ^ path ^ " Depends: " ^ (match path with
          | "index.js" -> "a b"
          | "a.js" -> "b c"
          | "c.js" -> ""
          | _ -> ""
        )
      | Resolve {jsdep; _} -> begin
          match jsdep with
          | "a" -> "a.js"
          | "b" -> "b.js"
          | "c" -> "c.js"
          | other -> failwith ("Module cannot be resolved: " ^ other)
        end
      | Compile (path, source) -> begin
          match source |> String.split_on_char ':' |> List.rev with
          | [] -> failwith "Should not happen"
          | deps :: _ ->
            let deps =
              deps
              |> String.trim
              |> String.split_on_char ' '
              |> List.filter (Pervasives.(<>) "")
            in
            (* TODO: compile should get the filename as well*)
            ({id = "id+" ^ path; source; path}, deps)
        end
      | Pack modules ->
        modules
        |> List.map (fun m -> m.source)
        |> String.concat "\n---\n"
  end


  include Build(Action)(Cache)

  (** Build primitives. *)

  (* Wrap read action as build family, allows to cache read results by id. *)
  let readSource = fun path ->
    define path (execute (Action.read path))

  (* Build which collects all module dependencies in a list. *)
  let rec buildGraph =
    fun compileModule path ->
      let%bind jsmodule, deps = compileModule path in

      (** Build which caches resolution dependencies on a jsmodule we are
       * resolving jsdep from and jsdep itself. *)
      let resolveDep = fun dep ->
          define dep (
              execute (Action.resolve ~context:jsmodule dep)
            )
      in

      define jsmodule.id (
        let%bind jsmodules =
          many (
            let f dep =
              let%bind path = resolveDep dep in
              buildGraph compileModule path
            in
            List.map f deps
          )
        in

        return (jsmodule :: (List.flatten jsmodules))
      )

  let gatherDeps depRequests =
    depRequests

  let resolve : jsmodule -> path -> path t = fun jsmodule dep ->
    define dep (
      execute (Action.resolve ~context:jsmodule dep)
    )

  let readCompileResolve: path -> (jsmodule * path list) t  = fun path ->
    define path (
      let%bind source = execute (Action.read path) in
      let%bind jsmodule, depRequests = define "ast" (execute (Action.compile path source)) in
      pair (return jsmodule) (many(List.map (resolve jsmodule) depRequests))
    )

  let buildGraph2 = fun entry ->
    let rec buildGraph' seen entry =
      let%bind jsmodule, resolvedDeps = readCompileResolve entry in
      let%bind jsmodules =
        many(
          resolvedDeps
          |> List.filter(fun dep -> not(StringSet.mem dep seen))
          |> List.map (buildGraph' (StringSet.add jsmodule.path seen))
        )
      in
      return (jsmodule :: (List.flatten jsmodules))
    in
    buildGraph' StringSet.empty entry


  (* Fastpack entry point / build which packs a bundle given an entry point. *)
  let main : path -> jsbundle t = fun entry ->
    define "bundle" (
      let%bind graph = define "graph" (buildGraph2 entry) in
      define "output.js" (execute (Action.pack graph))
    )
    (* (1* Build which caches module compilation. *1) *)
    (* let compileModule  = fun path -> *)
    (*   let%bind jssource = readSource path in *)
    (*   execute (Action.compile path jssource) *)
    (* in *)
    (* let%bind graph = buildGraph compileModule entry in *)
    (* execute (Action.pack graph) *)

end

let () =
  let bundle = Fastpack.(run (main "index.js")) in
  print_endline bundle