-- Annotation of normal primitives

module Yhc.Core.AnnotatePrims (
  buildPrimSpecMap
 ,normPrimSpecMap
 ,CoreStrictness (..)
 ,CoreTypeSig (..)
 ,buildPrimAnno
 ,buildNormPrimAnno
 ) where

import Data.Maybe
import Yhc.Core
import Yhc.Core.PrimAnnoRaw
import Yhc.Core.Annotation
import qualified Data.Map as M

-- Annotation key for functions and primitives

fparity (p@CorePrim {}) = corePrimArity p
fparity (f@CoreFunc {}) = length $ coreFuncArgs f

instance CoreAnnotable CoreFunc where
  toAnnotationKey (p@CorePrim {}) = "primitive_" ++ 
                                    coreFuncName p ++ "/" ++ 
                                    show (corePrimArity p)

  toAnnotationKey (f@CoreFunc {}) = "function_" ++
                                    coreFuncName f ++ "/" ++
                                    show (length $ coreFuncArgs f)

-- |Build a map of primitive specifications given the list of
-- primitive description records. This as well may be used by frontends.

buildPrimSpecMap :: [[String]] -> M.Map String [String]

buildPrimSpecMap pspc = M.fromList $ map bpsm pspc where
  bpsm (h:t) = (h, t)

-- |Specifications map of normal primitives

normPrimSpecMap :: M.Map String [String]

normPrimSpecMap = buildPrimSpecMap rawPrimAnno

-- Core function/primitive strictness is a list of Bools. True
-- at certain position means that a function or a primitive is strict
-- on the corresponding argument.

newtype CoreStrictness = CoreStrictness [Bool]

instance CoreProperty CoreStrictness where
  toAnnString (CoreStrictness bsct) = map (\b -> if b then 'T' else 'F') bsct
  fromAnnString s = mapM (\c -> case c of
                     'F' -> return False
                     'T' -> return True
                     _   -> fail $ "invalid strictness annotation: " ++ s) s
                >>= return . CoreStrictness

-- Core function/primitive type signature is a string containing Haskell
-- type expression.

newtype CoreTypeSig = CoreTypeSig String

instance CoreProperty CoreTypeSig where
  toAnnString (CoreTypeSig s) = s
  fromAnnString = return . CoreTypeSig

-- |Given the linked Core, build annotations for all primitives defined
-- that belong to the given set of primitives, that is, their names
-- are member keys of the given primitives specification map.

buildPrimAnno :: M.Map String [String] -> Core -> CoreAnnotations

buildPrimAnno mps core = ba M.empty (coreFuncs core) mps where
  ba am [] _ = am
  ba am (p:ps) mps | coreFuncName p `M.member` mps = 
    ba am'' ps mps where
      bsct art "All" = replicate art True
      bsct art "None" = replicate art False
      bsct art s = take art $ map ('T' ==) s ++ repeat False
      (use:descr:impl:arity:strct:tsig:_) = fromJust $ M.lookup (coreFuncName p) mps
      am'  = addAnnotation p ("Strictness", CoreStrictness (bsct (read arity) strct)) am
      am'' = addAnnotation p ("Type", CoreTypeSig tsig) am'
  ba am (_:ps) mps = ba am ps mps

-- |Given the linked Core, build annotations for all normal primitives
-- that is, belonging to the 'normPrimSpecMap'.

buildNormPrimAnno :: Core -> CoreAnnotations

buildNormPrimAnno = buildPrimAnno normPrimSpecMap