-- Annotation framework for Yhc Core

module Yhc.Core.Annotation (
  CoreAnnotations,
  CoreAnnotable (..),
  CoreProperty (..),
  addAnnotation,
  getAnnotation,
  combineAnnotations) where

import Data.Maybe
import qualified Data.Map as M
import qualified Data.List as L
import Yhc.Core

-- |Annotations database: a two-level map to hold property mappings 
-- for each of annotated objects.

type CoreAnnotations = M.Map String (M.Map String String)

-- |For each annotable object, unique key should be generated, to be used
-- with the top level map in the annotations database.

class CoreAnnotable a where
  toAnnotationKey :: a -> String

-- |For each property, an encoding (to String) and decoding (from String)
-- functions should be defined.

class CoreProperty p where
  toAnnString :: p -> String    -- ^arbitrary property to a string to store
  fromAnnString :: (Monad m) => String -> m p -- ^from stored string to property value or fail

-- |Given an annotable object, append a property with given name and value
-- to the existing annotations database.

addAnnotation :: (CoreAnnotable a, CoreProperty p) 
              => a -- ^annotable object
              -> (String, p) -- ^pair of name and value
              -> CoreAnnotations -- ^existing annotations database
              -> CoreAnnotations -- ^updated annotations database

addAnnotation a (n, p) b =
  let ak = toAnnotationKey a                   -- annotation key to lookup at top level
      aa = M.findWithDefault M.empty ak b      -- second level map or empty map if new
      ps = toAnnString p                       -- stringify the property value
      aa' = M.insert n ps aa                   -- insert/replace property at the second level
  in  M.insert ak aa' b                        -- insert/replace properties at the top level
  

getAnnotation :: (CoreAnnotable a, CoreProperty p) 
              => a -- ^annotable object
              -> String -- ^property name
              -> CoreAnnotations -- ^annotations database
              -> Maybe p -- ^returned value or nothing

getAnnotation a n b = do                       -- as we are in a monad...
  let ak = toAnnotationKey a                   -- annotation key to lookup at top level
  aa <- M.lookup ak b                          -- second level map or fail right here
  ps <- M.lookup n aa                          -- property value or fail right here
  p <- fromAnnString ps                        -- decode property value or fail rught here
  return p


-- |Given the two annotation sets, combine them into one. If the same object 
-- is annotated in both sets, annotations are combines for such object, 
-- and left annotations take precedence.

combineAnnotations :: CoreAnnotations -> CoreAnnotations -> CoreAnnotations

combineAnnotations l r = 
  let flkup = flip M.lookup
      lkeys = M.keys l
      rkeys = M.keys r
      ikeys = L.intersect lkeys rkeys
      lnodup = lkeys L.\\ ikeys
      rnodup = rkeys L.\\ ikeys
      lanno = concat $ map (maybeToList . flkup l) lnodup
      ranno = concat $ map (maybeToList . flkup r) rnodup
      ilanno = concat $ map (maybeToList . flkup l) ikeys
      iranno = concat $ map (maybeToList . flkup r) ikeys
      imps = zipWith M.union ilanno iranno
  in  M.fromList (zip lnodup lanno ++ zip rnodup ranno ++ zip ikeys imps)