{-# LANGUAGE TemplateHaskell, QuasiQuotes, PatternGuards, DoAndIfThenElse #-}

module Data.TrieMap.Representation.TH (genRepr) where

import Data.TrieMap.Rep.TH
import Data.TrieMap.Rep
import Data.TrieMap.Regular.Base
import Data.TrieMap.Key
import Data.TrieMap.Rep.Instances
import Language.Haskell.TH
import Language.Haskell.TH.ExpandSyns
import Control.Arrow
import Control.Monad

data ToRepCase = ToRepCase [Pat] Exp
data FromRepCase = FromRepCase Pat [Exp]
type ToRep = [ToRepCase]
type FromRep = [FromRepCase]

type Representation = (Type, ToRep, FromRep)

-- | Given the name of a type constructor, automatically generates an efficient 'Repr' instance.  /Warning/: Generalized tries do not work for "infinitely complicated types," for example, a type-system construction of the natural numbers.
-- In these cases, a context reduction stack overflow will occur at compile time when you use the 'TKey' instance for that type.
genRepr :: Name -> Q [Dec]
genRepr tycon = do
	TyConI dec <- reify tycon
	case dec of
		DataD _ _ tyvars cons _ -> do
			conReprs <- mapM conRepr cons
			return (decsForRepr (foldl AppT (ConT tycon) (map tyVarBndrType tyvars)) (foldr1 union conReprs))
		NewtypeD _ _ tyvars con _ -> do
			theConRepr <- conRepr con
			return (decsForRepr (foldl AppT (ConT tycon) (map tyVarBndrType tyvars)) theConRepr)

tyVarBndrType :: TyVarBndr -> Type
tyVarBndrType (PlainTV tyvar) = VarT tyvar
tyVarBndrType (KindedTV tyvar _) = VarT tyvar

decsForRepr :: Type -> Representation -> [Dec]
decsForRepr t (tRep, toR, fromR) = [
		TySynInstD ''Rep [t] tRep,
		InstanceD [] (ConT ''Repr `AppT` t)
			[FunD 'toRep
				[Clause pats (NormalB e) [] | ToRepCase pats e <- toR],
			 FunD 'fromRep
				[Clause [pat] (NormalB e) [] | FromRepCase pat [e] <- fromR]]]

decompose :: Type -> (Type, [Type])
decompose (tyfun `AppT` ty) = case decompose tyfun of
	(tyfun, tys)	-> (tyfun, tys ++ [ty])
decompose ty = (ty, [])

type ReprM = Q

conRepr :: Con -> ReprM Representation
conRepr (RecC con args) = conRepr (NormalC con [(strict, typ) | (_, strict, typ) <- args])
conRepr (InfixC t1 con t2) = conRepr (NormalC con [t1, t2])
conRepr (NormalC con []) = return $ conify con unit
conRepr (NormalC con args) = do
	argCons <- mapM (typeRepr . snd) args
	return (conify con (foldr1 prod argCons))

typeRepr :: Type -> ReprM Representation
typeRepr t00 = expandSyns t00 >>= \ t0 -> case decompose t0 of
	(ListT, [t])	-> do
		(tRep, toR, fromR) <- typeRepr t
		xs <- newName "elems"
		x <- newName "el"
		xsRep <- newName "elemReps"
		xRep <- newName "elemRep"
		return (ListT `AppT` tRep,
			[ToRepCase [VarP xs] 
				(CompE [BindS (VarP x) (VarE xs),
					NoBindS (CaseE (VarE x) [Match pat (NormalB e) [] | ToRepCase [pat] e <- toR])])],
			[FromRepCase (VarP xsRep)
				[CompE [BindS (VarP xRep) (VarE xsRep),
					NoBindS (CaseE (VarE xRep) [Match pat (NormalB e) [] | FromRepCase pat [e] <- fromR])]]])
	(TupleT 0, _)	-> return unit
	(TupleT n, ts)	-> do
		reps <- mapM typeRepr ts
		let (tRep, toR, fromR) = foldr1 prod reps
		return (tRep, [ToRepCase [TupP pats] e | ToRepCase pats e <- toR], [FromRepCase pat [TupE es] | FromRepCase pat es <- fromR])
	(ConT con, ts)
		| con == ''()	-> return unit
		| con == ''Either, [tL, tR] <- ts
			-> do	(tRepL, lToR, lFromR) <- typeRepr tL
				(tRepR, rToR, rFromR) <- typeRepr tR
				return (ConT ''Either `AppT` tRepL `AppT` tRepR,
					[ToRepCase [ConP 'Left pats] (ConE 'Left `AppE` e) | ToRepCase pats e <- lToR] ++
						[ToRepCase [ConP 'Right pats] (ConE 'Right `AppE` e) | ToRepCase pats e <- rToR],
					[FromRepCase (ConP 'Left [pat]) [ConE 'Left `AppE` e] | FromRepCase pat [e] <- lFromR] ++
						[FromRepCase (ConP 'Right [pat]) [ConE 'Right `AppE` e] | FromRepCase pat [e] <- rFromR])
		| con == ''Maybe, [t] <- ts
			-> do	(tRep, toR, fromR) <- typeRepr t
				return (ConT ''Either `AppT` TupleT 0 `AppT` tRep,
					[ToRepCase [ConP 'Nothing []] (ConE 'Left `AppE` TupE [])] ++
						[ToRepCase [ConP 'Just pats] (ConE 'Right `AppE` e) | ToRepCase pats e <- toR],
					[FromRepCase (RecP 'Left []) [ConE 'Nothing]] ++
						[FromRepCase (ConP 'Right [pat]) [ConE 'Just `AppE` e] | FromRepCase pat [e] <- fromR])
		| otherwise -> do
					ClassI _ instances <- reify ''Repr
					let knowns = [tycon | ClassInstance{ci_tys = [ConT tycon]} <- instances]
					if con `elem` knowns && null ts then do
						arg <- newName "arg"
						argRep <- newName "argRep"
						return (ConT ''Rep `AppT` ConT con,
							[ToRepCase [VarP arg] (VarE 'toRep `AppE` VarE arg)],
							[FromRepCase (VarP argRep) [VarE 'fromRep `AppE` VarE argRep]])
					else recursiveRepr t0
	_	-> recursiveRepr t0

tyVarBndrName :: TyVarBndr -> Name
tyVarBndrName (PlainTV n) = n
tyVarBndrName (KindedTV n _) = n

recursiveRepr :: Type -> ReprM Representation
recursiveRepr t0 = do	-- TODO: handle type synonyms here
		x <- newName "arg"
		return (ConT ''Key `AppT` t0, 
			[ToRepCase [VarP x] (ConE 'Key `AppE` VarE x)],
			[FromRepCase (ConP 'Key [VarP x]) [VarE x]])

unit :: Representation
unit = (TupleT 0, [ToRepCase [] (TupE [])], [FromRepCase WildP []])

prod :: Representation -> Representation -> Representation
prod (t1, toRep1, fromRep1)
	(t2, toRep2, fromRep2) =
	(TupleT 2 `AppT` t1 `AppT` t2,
		do	ToRepCase pats1 out1 <- toRep1
			ToRepCase pats2 out2 <- toRep2
			return (ToRepCase (pats1 ++ pats2) (TupE [out1, out2])),
		do	FromRepCase pat1 out1 <- fromRep1
			FromRepCase pat2 out2 <- fromRep2
			return (FromRepCase (TupP [pat1, pat2]) (out1 ++ out2)))

conify :: Name -> Representation -> Representation
conify conName (t, toR, fromR) =
	(t, [ToRepCase [ConP conName args] e | ToRepCase args e <- toR], [FromRepCase p [foldl AppE (ConE conName) outs] | FromRepCase p outs <- fromR])

union :: Representation -> Representation -> Representation
union (t1, toRep1, fromRep1)
	(t2, toRep2, fromRep2) =
	(ConT ''Either `AppT` t1 `AppT` t2,
		[ToRepCase pats (ConE 'Left `AppE` e) | ToRepCase pats e <- toRep1] ++
		[ToRepCase pats (ConE 'Right `AppE` e) | ToRepCase pats e <- toRep2],
		[FromRepCase (ConP 'Left [pat]) es | FromRepCase pat es <- fromRep1] ++
		[FromRepCase (ConP 'Right [pat]) es | FromRepCase pat es <- fromRep2])

-- genRepInstance :: Type -> Representationesentation -> Q Dec
-- genInstance