(* Copyright (c) 2008, Adam Chlipala
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * - Redistributions of source code must retain the above copyright notice,
 *   this list of conditions and the following disclaimer.
 * - Redistributions in binary form must reproduce the above copyright notice,
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 * - The names of contributors may not be used to endorse or promote products
 *   derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *)

structure Tag :> TAG = struct

open Core

structure U = CoreUtil
structure E = CoreEnv

structure IM = IntBinaryMap

fun kind (k, s) = (k, s)
fun con (c, s) = (c, s)

fun exp (e, s) =
    case e of
        EApp (
        (EApp (
         (EApp (
          (ECApp (
           (ECApp (
            (ECApp (
             (ECApp (
              (EFfi ("Basis", "tag"),
               loc), given), _), absent), _), outer), _), inner), _),
          attrs), _),
         tag), _),
        xml) =>
        (case attrs of
             (ERecord xets, _) =>
             let
                 val (xets, s) =
                     ListUtil.foldlMap (fn ((x, e, t), (count, tags, newTags)) =>
                                           case x of
                                               (CName "Link", _) =>
                                               let
                                                   fun unravel (e, _) =
                                                       case e of
                                                           ENamed n => (n, [])
                                                         | EApp (e1, e2) =>
                                                           let
                                                               val (n, es) = unravel e1
                                                           in
                                                               (n, es @ [e2])
                                                           end
                                                         | _ => (ErrorMsg.errorAt loc "Invalid link expression";
                                                                 (0, []))

                                                   val (f, args) = unravel e

                                                   val (cn, count, tags, newTags) =
                                                       case IM.find (tags, f) of
                                                           NONE =>
                                                           (count, count + 1, IM.insert (tags, f, count),
                                                            (f, count) :: newTags)
                                                         | SOME cn => (cn, count, tags, newTags)

                                                   val e = (EClosure (cn, args), loc)
                                                   val t = (CFfi ("Basis", "string"), loc)
                                               in
                                                   ((x, e, t),
                                                    (count, tags, newTags))
                                               end
                                             | _ => ((x, e, t), (count, tags, newTags)))
                     s xets
             in
                 (EApp (
                  (EApp (
                   (EApp (
                    (ECApp (
                     (ECApp (
                      (ECApp (
                       (ECApp (
                        (EFfi ("Basis", "tag"),
                         loc), given), loc), absent), loc), outer), loc), inner), loc),
                    (ERecord xets, loc)), loc),
                   tag), loc),
                  xml), s)
             end
           | _ => (ErrorMsg.errorAt loc "Attribute record is too complex";
                   (e, s)))

      | _ => (e, s)

fun decl (d, s) = (d, s)

fun tag file =
    let
        val count = foldl (fn ((d, _), count) =>
                              case d of
                                  DCon (_, n, _, _) => Int.max (n, count)
                                | DVal (_, n, _, _, _) => Int.max (n, count)
                                | DExport _ => count) 0 file

        fun doDecl (d as (d', loc), (env, count, tags)) =
            let
                val (d, (count, tags, newTags)) =
                    U.Decl.foldMap {kind = kind,
                                    con = con,
                                    exp = exp,
                                    decl = decl}
                                   (count, tags, []) d

                val env = E.declBinds env d

                val newDs = ListUtil.mapConcat
                                (fn (f, cn) =>
                                    let
                                        fun unravel (all as (t, _)) =
                                            case t of
                                                TFun (dom, ran) =>
                                                let
                                                    val (args, result) = unravel ran
                                                in
                                                    (dom :: args, result)
                                                end
                                              | _ => ([], all)

                                        val (fnam, t, _, tag) = E.lookupENamed env f
                                        val (args, result) = unravel t

                                        val (app, _) = foldl (fn (t, (app, n)) =>
                                                                 ((EApp (app, (ERel n, loc)), loc),
                                                                  n - 1))
                                                             ((ENamed f, loc), length args - 1) args
                                        val body = (EWrite app, loc)
                                        val unit = (TRecord (CRecord ((KType, loc), []), loc), loc)
                                        val (abs, _, t) = foldr (fn (t, (abs, n, rest)) =>
                                                                    ((EAbs ("x" ^ Int.toString n,
                                                                            t,
                                                                            rest,
                                                                            abs), loc),
                                                                     n + 1,
                                                                     (TFun (t, rest), loc)))
                                                                (body, 0, unit) args
                                    in
                                        [(DVal ("wrap_" ^ fnam, cn, t, abs, tag), loc),
                                         (DExport cn, loc)]
                                    end) newTags
            in
                (newDs @ [d], (env, count, tags))
            end

        val (file, _) = ListUtil.foldlMapConcat doDecl (CoreEnv.empty, count, IM.empty) file
    in
        file
    end

end