-- | Monomorphises the Simple language, given some initial activated
--   records. The idea is that these activation records come from
--   a monomorphised conjecture.
{-# LANGUAGE RecordWildCards, ScopedTypeVariables, PatternGuards #-}
module HipSpec.Lang.Monomorphise (monoClauses, IdInst(..),Lemma(..),toList,Records) where

import Prelude hiding (lookup)
import HipSpec.Lang.PolyFOL
import qualified Data.Set as S
import Data.Set (Set)
import qualified Data.Map as M
import Data.Map (Map)

import Data.Monoid
import Data.Maybe
import Data.List (partition)
import Control.Arrow
import Control.Monad

type Records a b = Map (Trigger a) (Set [Type a b])

empty :: Records a b
empty = M.empty

member :: (Ord a,Ord b) => (Trigger a,[Type a b]) -> Records a b -> Bool
(t,ts) `member` r = maybe False (ts `S.member`) (M.lookup t r)

(\\) :: (Ord a,Ord b) => Records a b -> Records a b -> Records a b
(\\) = M.differenceWith $ \ s1 s2 -> do
    let s' = s1 S.\\ s2
    guard (not (S.null s'))
    return s'

union :: (Ord a,Ord b) => Records a b -> Records a b -> Records a b
union = M.unionWith S.union

unions :: (Ord a,Ord b) => [Records a b] -> Records a b
unions = M.unionsWith S.union

fromList :: (Ord a,Ord b) => [(Trigger a,[Type a b])] -> Records a b
fromList = M.fromListWith S.union . map (second S.singleton)

toList :: (Ord a,Ord b) => Records a b -> [(Trigger a,[Type a b])]
toList r = [ (t,ts) | (t,s) <- M.toList r, ts <- S.toList s ]

minView :: (Ord a,Ord b) => Records a b -> Maybe ((Trigger a,[Type a b]),Records a b)
minView m = case M.minViewWithKey m of
    Just ((t,s),m') -> case S.minView s of
        Just (ts,s') -> Just ((t,ts),M.insert t s' m')
        Nothing -> minView m'
    Nothing -> Nothing

insert :: (Ord a,Ord b) => Trigger a -> [Type a b] -> Records a b -> Records a b
insert i t = M.insertWith S.union i (S.singleton t)

lookup :: (Ord a,Ord b) => Trigger a -> Records a b -> [[Type a b]]
lookup t r = maybe [] S.toList (M.lookup t r)

-- | Get the records from a clause
clauseRecs :: (Ord a,Ord b) => Clause a b -> Records a b
clauseRecs cl = fromList $
    [ (TySymb tc,ts) | TyCon tc ts  <- clTyUniv cl ] ++
    [ (Symb f,ts)    | Apply f ts _ <- clTmUniv cl ]

instClauseWith :: forall a b . Eq b => Clause a b -> [(b,Type a b)] -> Clause a b
instClauseWith cl su = case cl of
    Clause nm _trg ty _tvs fm -> Clause nm [] ty [] (fmInsts su fm)
    _ -> cl -- or error!

instClause :: forall a b . Eq b => Clause a b -> [Type a b] -> Clause a b
instClause cl ts = case cl of
    Clause{..} -> instClauseWith cl (zip ty_vars ts)
    _ -> cl

type InstMap a b = Map b (Type a b)

data Lemma a b = Lemma
    { lm_cl :: Clause a b
    -- ^ the clause
    , lm_act :: [(Trigger a,[Type a b])] -- Records a b
    -- ^ The instantiation records it requires
    , lm_eff :: Records a b
    -- ^ The effect of instantiating it
    , lm_inst :: [InstMap a b]
    -- ^ The types it has already been instantiated at
    }
  deriving Show

-- | Stage one monomorphisation
--   1st component: monomorphically applied clauses
--   2nd component: type signatures to be monomorphised
monoClauses1
    :: forall a b . (Ord a,Ord b) => [Clause a b]
    -> (([Clause a b],[Clause a b]),([Lemma a b],Records a b))
monoClauses1 cls0 = ((source ++ defs ++ lemma_cls,sigs),fin)
  where
    trg_pred p Clause{..} = p cl_ty_triggers
    trg_pred _ _          = False
    (source,other)        = partition (trg_pred (Source `elem`)) cls0
    (defs_poly,rest)      = partition (trg_pred (not . null))    other
    (lem_cls,sigs)        = partition (trg_pred (const True))    rest

    mono = mkMono defs_poly

    (defs,def_irs) = mono empty (unions (map clauseRecs source))

    lemmas :: [Lemma a b]
    lemmas =
        [ Lemma cl (toList recs) (snd (mono empty recs)) []
        | cl <- lem_cls
        , let recs = clauseRecs cl
        ]

    go :: Int -> Bool -> [Lemma a b] -> [Lemma a b] -> Records a b -> ([Clause a b],([Lemma a b],Records a b))
    go _rounds False [] acc irs = ([],(acc,irs))
    go rounds  True  [] acc irs = go (rounds+1) False acc [] irs
    go rounds  b (l:ls) acc irs = case new of
        []        -> go rounds b ls (l:acc) irs
        (l',cl):_ ->
            let (cls,irs') = mono irs (clauseRecs cl)
            in  first (\ c -> cl:cls ++ c) (go rounds True ls (l':acc) irs')
      where
        new = catMaybes
            [ instLemma (rounds < 2) l im irs
            | im <- possibleInsts l irs
            ]

    (lemma_cls,fin) = go 0 False lemmas [] def_irs

instLemma :: (Ord a,Ord b)
          => Bool        {- ^ allow new type instantiations -}
          -> Lemma a b   {- ^ lemma to maybe instantiate -}
          -> InstMap a b {- ^ type to instantiate it at -}
          -> Records a b {- ^ instantiated records so far -}
          -> Maybe (Lemma a b,Clause a b)
                         {- Just (l,r,c) :
                                l: Lemma with new info
                                c: The clause that got instantiated -}
instLemma allow_new l@Lemma{..} im rs
    -- Instantiate lemmas as long as they don't trigger new
    -- types to be instantiated
    | (allow_new || and
        [ (t,map (tySubsts (M.toList im)) ts) `member` rs
        | (t,ts) <- takeWhile (isTySymb . fst) (toList lm_eff)
        ])
      -- and we haven't instantiated it before
      && im `notElem` lm_inst
      -- and all type variables are assigned
      && all (`M.member` im) (ty_vars lm_cl)

    = Just
        ( l { lm_inst = im : lm_inst }
        , instClauseWith lm_cl (M.toList im)
        )

    | otherwise = Nothing

-- | Possible instantiations of a lemma given some instantiation records
possibleInsts :: (Ord a,Ord b) => Lemma a b -> Records a b -> [InstMap a b]
possibleInsts Lemma{..} rs = go lm_act
  where
    go []            = [M.empty]
    go ((trg,ts):xs) =
        [ m1 `M.union` m2
        | m1 <- M.empty : mapMaybe (tyInsts ts) (lookup trg rs)
        , m2 <- go xs
        , compatible m1 m2
        ]

tyInst :: (Ord a,Ord b)
       => Type a b {- ^ with type variables -}
       -> Type a b {- ^ without -}
       -> Maybe (InstMap a b)
tyInst (TyVar x)    t                     = Just (M.singleton x t)
tyInst (TyCon u us) (TyCon v vs) | u == v = tyInsts us vs
tyInst TType        TType                 = Just M.empty
tyInst Integer      Integer               = Just M.empty
tyInst _            _                     = Nothing

tyInsts :: (Ord a,Ord b) => [Type a b] -> [Type a b] -> Maybe (InstMap a b)
tyInsts (u:us) (v:vs) = do
    m1 <- tyInst u v
    m2 <- tyInsts us vs
    guard (compatible m1 m2)
    return (m1 `M.union` m2)
tyInsts [] [] = Just M.empty
tyInsts _  _  = Nothing

compatible :: (Ord a,Ord b) => InstMap a b -> InstMap a b -> Bool
compatible a b = oneway a b && oneway b a
  where
    oneway p q = and [ maybe True (== v) (M.lookup k q) | (k,v) <- M.toList p ]

mkMono :: forall a b . (Ord a,Ord b)
       => [Clause a b]
       -> Records a b                -- ^ Initial, already instantiated records
       -> Records a b                -- ^ New record queue
       -> ([Clause a b],Records a b) -- ^ Instantiated clauses and final records
mkMono cls0 irs0 q0 = go irs0 (q0 \\ irs0)
  where
    -- Since some clauses are activated by either of many triggers,
    -- we let one of them instantiate the clause, and the other ones just
    -- retrigger with the major pattern
    trg_targets :: Map (Trigger a) ([Clause a b],[Trigger a])
    trg_targets = M.fromListWith mappend $ concat $
        [ (trg,([cl],rtrs)) : [ (rtr,([],filter (rtr /=) trgs)) | rtr <- rtrs ]
        | cl <- cls0
        , let trgs@(trg:rtrs) = cl_ty_triggers cl
        ]

    go :: Records a b
       -> Records a b
       -- ^ (invariant: no overlap)
       -> ([Clause a b],Records a b)
    go irs q = case minView q of
        Nothing -> ([],irs)
        Just ((trg,ts),q') -> case M.lookup trg trg_targets of
            Nothing            -> go irs q'
            Just (cls,retrigs) ->
                let cls' = map (`instClause` ts) cls
                    rtr  = fromList (zip retrigs (repeat ts))
                    irs' = insert trg ts irs
                    recs = (unions (map clauseRecs cls') `union` rtr) \\ irs'
                in  first (cls' ++) (go irs' (q' `union` recs))

data IdInst a b = IdInst a [Type a b]
  deriving (Eq,Ord,Show)

-- | Second pass monomorphisation: remove all type applications and change
--   the identifier names instead
monoClauses2 :: (Ord a,Ord b) => ([Clause a b],[Clause a b]) -> [Clause (IdInst a b) b]
monoClauses2 (cls,sigs) = map (`SortSig` 0) (S.toList sorts) ++ ty_sigs ++ cls'
  where
    cls' =
        [ Clause
            { cl_name        = cl_name cl
            , cl_ty_triggers = []
            , cl_type        = cl_type cl
            , ty_vars        = []
            , cl_formula     = fmMod tyCon apply (cl_formula cl)
            }
        | cl <- cls
        ]

    tyCon f ts = IdInst f ts
    apply f ts = Apply (IdInst f ts) []

    (sorts1,ty_apps) = clsDeps cls'


    sig_map = M.fromList [ (f,(tvs,args,res)) | TypeSig f tvs args res <- sigs ]

    ty_sigs =
       [ TypeSig f' [] args' res'
       | f'@(IdInst f ts) <- S.toList ty_apps
       , Just (tvs,args,res) <- [M.lookup f sig_map]
       , let su = tyMod tyCon . tySubsts (zip tvs ts)
             args' = map su args
             res'  = su res
       ]

    (sorts2,_) = clsDeps ty_sigs

    sorts = sorts1 `S.union` sorts2

-- | Monomorphise clauses
monoClauses :: (Ord a,Ord b) => [Clause a b] -> ([Clause (IdInst a b) b],([Lemma a b],Records a b))
monoClauses = first monoClauses2 . monoClauses1

