Haskell Code by HsColour

{-# OPTIONS_GHC -fno-warn-missing-methods #-}

module Data.DeriveGuess(DataName(..), tup1, guess) where

import Language.Haskell.TH.All
import Data.Generics
import Data.List
import Data.Char
import Data.Maybe


data DataName a = CtorZero
 | CtorOne a
 | CtorTwo a a
 | CtorTwo' a a

ctorNames = ["CtorZero","CtorOne","CtorTwo","CtorTwo'"]

-- | Guess an instantiator from a sample instance.
guess :: (String, Q [Dec]) -> IO ()
guess (name,x) = runQ x >>= putStr . (++) line0. widthify . (++) line1 . guessStr . unQ
 where
 line0 = "make" ++ name ++ " :: Derivation\n" ++
 "make" ++ name ++ " = derivation " ++ lname ++ "' \"" ++ name ++ "\"\n"
 line1 = lname ++ "' dat = "
 lname = toLower (head name) : tail name

-- | A fake constructor for the unary tuple. Helps 'guess' to see
-- patterns in progressions of differently sized tuples.
tup1 = id

-- | Chop and mangle a String representing Haskell code so that it
-- fits in 80 columns, without regard for prettiness.
widthify :: String -> String
widthify xs = g 80 (f xs)
 where
 g n (x:xs) | n - length x <= 0 = "\n " ++ g 76 ([x|x/=" "] ++ xs)
 | otherwise = x ++ g (n - length x) xs
 g n [] = "\n"


 f (x:xs) | isSpace x = " " : f (dropWhile isSpace xs)
 f x = case lex x of
 [("","")] -> []
 -- \\ must not occur at the end of a line (CPP restrictions)
 [("\\",y)] -> let a:b = f y in ('\\':a) : b
 [(x,y)] -> x : f y

-- | Process a tree produced by a quasiquote, stripping name
-- uniquifiers and changing applications and tuplings into a standard
-- form.
unQ :: [Dec] -> [Dec]
unQ x = everywhere (mkT g) $ everywhere fAny $ map normData x
 where
 fAny :: (Typeable a, Data a) => a -> a
 fAny = mkT fE `extT` fP
 fE (VarE x) = VarE (f x); fE x = x
 fP (VarP x) = VarP (f x); fP x = x

 -- | Remove_0 evil_1 ghc_2 name_3 uniquifiers_4
 f :: Name -> Name
 f name = if not ("_" `isSuffixOf` s) && match s then mkName $ dropUnder s else name
 where
 s = show name
 match = isPrefixOf "_" . dropWhile isDigit . reverse

 -- | Turn infix applications into prefix, and normalise
 -- tuples.
 g :: Exp -> Exp
 g (InfixE (Just x) y (Just z)) = AppE (AppE y x) z
 g (AppE (VarE tup) x) | show tup == "tup1" = TupE [x]
 g (ConE unit) | show unit == "()" = TupE []
 g x = x


-- | Drop the first _ and everything after it; used to trim GHC name
-- uniques.
dropUnder = reverse . drop 1 . dropWhile (/= '_') . reverse

list x = "[" ++ concat (intersperse "," x) ++ "]"

unwordsb x = "(" ++ unwords x ++ ")"

fst3 (a,b,c) = a
snd3 (a,b,c) = b
thd3 (a,b,c) = c

snub x = nub $ sort x

arityToCtors x = x : [3 | x == 2]
ctorToArity x = if x == 3 then 2 else x

ctorArityEnv (Ctor i) = ctorToArity i

on op get a b = op (get a) (get b)


-- imagine the following environment table:
{-
[("CtorZero",0,0,[])
,("CtorOne" ,1,1,[1])
,("CtorTwo" ,2,2,[1,2])
,("CtorTwo'",2,3,[1,2])
]
-}

data Env = None | Ctor Int | Field Int
 deriving (Show,Eq)

isField (Field _) = True; isField _ = False
isCtor (Ctor _) = True; isCtor _ = False

fromField (Field i) = i
fromCtor (Ctor i) = i

fromEnv (Field i) = i
fromEnv (Ctor i) = i



-- Show t and Ppr t for better error messages
-- Eq t required for hypothesis testing
class (Ppr t, Eq t, Show t) => Guess t where
 -- invariant: all answers must be correct for this example
 -- will never be given a different type of environment
 guessEnv :: t -> [(Env, Env -> t, String)]
 
guessStr :: Guess t => t -> String
guessStr t = case [s | (None,_,s) <- guessEnv t] of
 [] -> error $ "\n\nNo hypothesis for:\n" ++ show t ++
 "\n\nPretty version:\n" ++ show (ppr t)
 (x:xs) -> x


checkGuess :: (Ppr t, Eq t, Show t) => t -> [(Env, Env -> t, String)] -> [(Env, Env -> t, String)]
checkGuess t xs = map f xs
 where
 f o@(env,gen,str) | t == gen env = o
 | otherwise = error $ unlines ["checkGuess failed:"
 ,"INPUT : " ++ show t
 ,"OUTPUT: " ++ show (gen env)
 ,"ENV : " ++ show env
 ,"HYP : " ++ str
 ]


guessEnvStr :: Guess t => t -> [(Env, Env -> t, String)]
guessEnvStr t = [(None, const t, guessStr t)]


guessPairStr :: (Guess a, Guess b) => String -> a -> b -> String
guessPairStr sjoin a b = sjoin ++ " " ++ guessStr a ++ " " ++ guessStr b

guessTripStr :: (Guess a, Guess b, Guess c) => String -> a -> b -> c -> String
guessTripStr sjoin a b c = unwords [sjoin, guessStr a, guessStr b, guessStr c]


joinEnvs :: [Env] -> Maybe Env
joinEnvs xs = if length ys > 1 then Nothing else Just $ head (ys ++ [None])
 where ys = filter (/= None) $ nub xs

guessOneEnv :: Guess a => (a -> t) -> String -> a -> [(Env, Env -> t, String)]
guessOneEnv fjoin sjoin x1 =
 [ (e1, \e -> fjoin (f1 e), unwordsb [sjoin,s1])
 | (e1,f1,s1) <- guessEnv x1]

-- to join two elements either they are the same env, or one has None
guessPairEnv :: (Guess a, Guess b) => (a -> b -> t) -> String -> a -> b -> [(Env, Env -> t, String)]
guessPairEnv fjoin sjoin x1 x2 =
 [ (env, \e -> fjoin (f1 e) (f2 e), unwordsb [sjoin,s1,s2])
 | (e1,f1,s1) <- guessEnv x1
 , (e2,f2,s2) <- guessEnv x2
 , Just env <- [joinEnvs [e1,e2]]]

guessTripEnv :: (Guess a, Guess b, Guess c) => (a -> b -> c -> t) -> String -> a -> b -> c -> [(Env, Env -> t, String)]
guessTripEnv fjoin sjoin x1 x2 x3 =
 [ (env, \e -> fjoin (f1 e) (f2 e) (f3 e), unwordsb [sjoin,s1,s2,s3])
 | (e1,f1,s1) <- guessEnv x1
 , (e2,f2,s2) <- guessEnv x2
 , (e3,f3,s3) <- guessEnv x3
 , Just env <- [joinEnvs [e1,e2,e3]]]



instance Guess a => Guess [a] where
 guessEnv os = concatMap f $ mapM guessEnv os
 where
 -- first try and induct based on the length of the list
 f xs | all (== None) (map fst3 xs) &&
 length xs == 2 &&
 length vals == 1
 = [(Ctor i, \e -> replicate (ctorArityEnv e) (head vals),
 "(replicate (ctorArity ctor) " ++ thd3 (head xs) ++ ")")
 | i <- [2,3]] ++
 [(None, \e -> map ($ e) gens, list strs)]
 where
 (envs,gens,strs) = unzip3 xs
 vals = nub $ zipWith ($) gens envs

 
 f xs | length es <= 1 = [(head (es ++ [None]), \e -> map ($ e) gens, list strs)]
 | otherwise = [(env,gen,"("++str++")")
 | env <- newEnvs, (gen,str) <- nubBy ((==) `on` snd) $ g xs]
 where
 (envs,gens,strs) = unzip3 xs
 es = nub $ filter (/= None) envs
 
 ctors = snub [i | Ctor i <- envs]
 fields = snub [i | Field i <- envs]
 maxField = maximum fields
 
 newEnvs = case ctors of
 [] -> map Ctor $ arityToCtors maxField
 _ | null fields -> [None]
 [x] | ctorToArity x == maxField -> [Ctor x]
 _ -> [] 
 
 ctorEnv = head newEnvs == None
 varName = if ctorEnv then "(ctorInd,ctor)" else "field"
 
 domain = if ctorEnv then [0..3] else [1..maxField]
 getDomain (Ctor i) = take 2 [1..i]
 getDomain None = [0..3]
 getDomain _ = []
 strDomain = if ctorEnv then "(zip [0..] (dataCtors dat))" else "[1..ctorArity ctor]"
 
 construct = if ctorEnv then Ctor else Field
 
 isNone x = x == None || (not ctorEnv && isCtor x)
 
 g :: Eq t => [(Env, Env -> t, String)] -> [(Env -> [t], String)]
 g [] = [(\e -> [], "[]")]
 g ((none,gn,st):xs) | isNone none =
 [(\e -> gn e : gen e, "[" ++ st ++ "]++" ++ str) | (gen,str) <- g xs]
 
 g xs = h id "id" xs ++ h reverse "reverse" xs

 h :: Eq t => ([Int] -> [Int]) -> String -> [(Env, Env -> t, String)] -> [(Env -> [t], String)]
 h fdir sdir xs
 | map construct (fdir domain) `isPrefixOf` map fst3 xs
 = [(\e -> map (fhyp . construct) (fdir $ getDomain e) ++ gen e
 ,"(map (\\" ++ varName ++ " -> " ++ shyp ++ ") (" ++ sdir ++ " " ++ strDomain ++ "))++" ++ str)
 | (fhyp,shyp) <- validHyp
 , (gen,str) <- g rest]
 where
 (now,rest) = splitAt (length domain) xs

 validHyp = filter (\hyp -> all (valid (fst hyp)) now) (map (\(a,b,c) -> (b,c)) now)
 valid hyp (e,gen,_) = hyp e == gen e

 h _ _ _ = []


guessType :: Type -> Type -> String
guessType principle x =
 if hasPrinciple then "(map (\\tdat -> " ++ disp x ++ ") (dataVars dat))" else "[" ++ disp x ++ "]"
 where
 hasPrinciple = f x
 where
 f x | x == principle = True
 f (AppT (ConT x) y) | show x == "DataName" = False
 f (AppT a b) = f a || f b
 f _ = False

 disp x | x == principle = "tdat"
 disp (AppT (ConT x) y) | show x == "DataName" = "(lK (dataName dat) (dataVars dat))"
 disp (AppT a b) = "(AppT " ++ disp a ++ " " ++ disp b ++ ")"
 disp (VarT a) = "(VarT (mkName " ++ show (show a) ++ "))"
 disp (ConT a) = "(ConT (mkName " ++ show (show a) ++ "))"



instance Guess Dec where
 guessEnv (InstanceD ctx typ inner) =
 [ (None, \e -> InstanceD ctx typ (gen e), prefix ++ str)
 | (None,gen,str) <- guessEnv inner]
 where
 principle = head (everything (++) ([] `mkQ` f) typ ++ [VarT $ mkName "?"])
 where
 f (AppT (ConT x) y) | show x == "DataName" = [y]
 f _ = []
 
 prefix = "InstanceD " ++
 "(concat (" ++ list (map (guessType principle) ctx) ++ ")) " ++
 "(head " ++ guessType principle typ ++ ")"

 guessEnv (FunD name claus) = guessPairEnv FunD "FunD" name claus
 guessEnv (ValD pat bod whr) = guessTripEnv ValD "ValD" pat bod whr
 
 guessEnv x = error $ show ("Guess Dec",x)


instance Guess Name where
 guessEnv name = if null guessCtor then guessRest else guessCtor
 where
 sname = show name
 (pre,end) = (init sname, last sname)
 
 guessCtor = [(Ctor i, \(Ctor e) -> mkName (pre ++ (ctorNames !! e))
 ,"(mkName (" ++ show pre ++ " ++ ctorName ctor))")
 | (i,nam) <- zip [0..] ctorNames, nam `isSuffixOf` sname
 , let pre = take (length sname - length nam) sname]

 guessRest = guessLast ++ guessDefault
 
 guessLast | isDigit end = [(e, \e -> mkName $ pre ++ show (g e)
 ,"(mkName (" ++ show pre ++ " ++ show " ++ s ++ "))")
 | (e,g,s) <- guessNum $ read [end]]
 | otherwise = []

 guessDefault = [(None,const name, "(mkName " ++ show sname ++ ")")
 | not (isDigit end) || pre `notElem` ["x","y","z"]]


guessNum :: Int -> [(Env, Env -> Int, String)]
guessNum i = [(Field i, fromField, "field") | i `elem` [1,2]] ++
 [(None, const 3, "(toInteger (length (dataCtors dat) - 1))") | i == 3] ++
 [(None, const 4, "(toInteger (length (dataCtors dat)))") | i == 4] ++
 [(Ctor i, fromCtor, "ctorInd") | i `elem` [0..3]] ++
 [(Ctor i, ctorArityEnv, "(ctorArity ctor)") | i `elem` [0..2]] ++
 [(Ctor 3, ctorArityEnv, "(ctorArity ctor)") | i == 2]



instance Guess Clause where
 guessEnv (Clause pat bod whr) = guessTripEnv Clause "Clause" pat bod whr


instance Guess Stmt where
 guessEnv (BindS x y) = guessPairEnv BindS "BindS" x y
 guessEnv (NoBindS x) = guessOneEnv NoBindS "NoBindS" x
 guessEnv x = error $ show ("Guess Stmt",x)


instance Guess Pat where
 guessEnv (VarP x) = guessOneEnv VarP "VarP" x
 guessEnv (ConP x xs) = guessPairEnv ConP "ConP" x xs
 guessEnv (WildP) = [(None, const WildP, "WildP")]
 guessEnv (TildeP x) = guessOneEnv TildeP "TildeP" x
 guessEnv (RecP x []) = guessOneEnv (flip RecP []) "(flip RecP [])" x
 guessEnv (LitP x) = guessOneEnv LitP "LitP" x
 guessEnv x = error $ show ("Guess Pat",x)


instance Guess Body where
 guessEnv (NormalB x) = guessOneEnv NormalB "NormalB" x
 guessEnv x = error $ show ("Guess Body",x)


instance Guess Exp where
 guessEnv (VarE x) = guessOneEnv VarE "VarE" x
 guessEnv (ConE x) = guessOneEnv ConE "ConE" x
 guessEnv (LitE x) = guessOneEnv LitE "LitE" x
 guessEnv (ListE x) = guessOneEnv ListE "ListE" x
 guessEnv (LamE x y) = guessPairEnv LamE "LamE" x y
 guessEnv (CompE x) = guessOneEnv CompE "CompE" x
 guessEnv (CaseE x y) = guessPairEnv CaseE "CaseE" x y
 guessEnv (TupE x) = guessOneEnv TupE "TupE" x
 guessEnv (RecConE x []) = guessOneEnv (flip RecConE []) "(flip RecConE [])" x
 guessEnv (CondE x y z) = guessTripEnv CondE "CondE" x y z
 guessEnv (DoE x) = guessOneEnv DoE "DoE" x

 guessEnv o@(AppE x y) = guessApply o ++ guessFold o ++ guessPairEnv AppE "AppE" x y
 
 guessEnv x = error $ show ("Guess Exp",x)


instance Guess Match where
 guessEnv (Match a b c) = guessTripEnv Match "Match" a b c


instance Guess Lit where
 guessEnv o@(IntegerL i) =
 [ (env, \e -> IntegerL $ toInteger $ gen e, "(IntegerL " ++ str ++ ")")
 | (env,gen,str) <- guessNum $ fromInteger i] ++
 [(None,const $ IntegerL i,"(IntegerL " ++ show i ++ ")")]
 
 guessEnv o@(StringL s) | s == "DataName" = [(None, const o, "(StringL (dataName dat))")]
 | otherwise = [(None, const o, "(StringL " ++ show s ++ ")")]
 
 guessEnv x = error $ show ("Guess Lit",x)


-- for when an expression is just an application
guessApply :: Exp -> [(Env, Env -> Exp, String)]
guessApply o | length args <= 1 = []
 | otherwise = guessPairEnv applyWith "applyWith" fn args
 where
 (fn,args) = list o
 
 list (AppE x y) = let (fn,args) = list x in (fn, args ++ [y])
 list x = (x, [])


-- for when an expression comes from folding
guessFold :: Exp -> [(Env, Env -> Exp, String)]
guessFold o@(AppE (AppE fn x) y) =
 f (with foldl1With) "foldl1With" (list True o) ++
 f (with foldr1With) "foldr1With" (list False o)
 where
 with fold join [] = VarE $ mkName "?"
 with fold join xs = fold join xs
 
 list b (AppE (AppE fn2 x) y) | fn == fn2 =
 if b then x : list b y else y : list b x
 list b x = [x]

 f ffold sfold lst
 | length lst <= 2 = []
 | otherwise = guessPairEnv ffold sfold fn lst

guessFold _ = []