Initial commit
[random-haskell-code.git] / RandomHaskell.hs
1 {-# LANGUAGE StandaloneDeriving, RankNTypes, ScopedTypeVariables #-}
2
3 import Language.Haskell.Exts          -- haskell-src-exts
4 import Data.Generics.Uniplate.Data    -- uniplate
5 import Data.List
6 import Data.Data
7 import System.Environment
8 import Control.Monad.Random
9 import qualified Data.Map.Strict as M
10 import Data.Functor
11 import Control.Monad
12 import Data.Maybe
13 import Unsafe.Coerce
14 import System.IO
15
16 deriving instance Ord ConstrRep
17
18 data MyCon = MyCon ConstrRep | Str String deriving (Eq, Ord)
19
20 type Key = (TypeRep, ShortCtxt)
21 type ConOccs = [(Key, MyCon)]
22 type ConMap = M.Map MyCon Int
23 type TyMap a = M.Map Key a
24
25 type Ctxt = [(TypeRep, MyCon, Int)]
26 type ShortCtxt = Ctxt
27
28 -- This is an interesting knob to turn. Larger means more realistic code, but
29 -- also more memory requirement and more code that should be read
30 shortenCtxt :: Ctxt -> ShortCtxt
31 shortenCtxt = take 3
32
33 extendCtxt :: Data a => a -> MyCon -> Int -> Ctxt -> Ctxt
34 -- Ignore one-constructor, one field data types
35 extendCtxt x (MyCon c) i ctxt
36     | AlgRep [c'] <- dataTypeRep (dataTypeOf x)
37     , [_] <- constrFields c'
38     = ctxt
39 extendCtxt x mc i ctxt = (typeOf x,mc,i) : ctxt
40
41 genListLength :: Data a => a -> Maybe Int
42 genListLength x = do
43   let t = typeOf x
44   (tc,[a]) <- return $ splitTyConApp t
45   guard (tc == fst (splitTyConApp (typeOf "hi")))
46   let l = unsafeCoerce x :: [()]
47   return (length l)
48
49
50 getConOccs :: Data a => Ctxt -> a -> ConOccs
51 getConOccs ctxt x 
52     | Just s <- cast x
53     = [ (key, Str s) ]
54     | otherwise
55     = let mc = MyCon (constrRep (toConstr x))
56           ctxt' i = extendCtxt x mc i ctxt
57       in (key, mc) : 
58          concat (zipWith (\f i -> (f i))
59                     (gmapQ (\y i -> getConOccs (ctxt' i) y) x)
60                     [0..])
61   where
62     key = (t, shortenCtxt ctxt)
63     t = typeOf x
64
65 singletonConMap :: MyCon -> ConMap
66 singletonConMap c = M.singleton c 1
67
68 conMapAdd :: ConMap -> ConMap -> ConMap
69 conMapAdd = M.unionWith (+)
70
71 toTypeMap :: ConOccs -> TyMap ConMap
72 toTypeMap = foldl (\m (t,c) -> M.insertWith conMapAdd t (singletonConMap c) m) M.empty
73
74 selectRandom :: Double -> M.Map r Int -> r
75 selectRandom r m = go i (M.toList m)
76   where
77   go i ((x, n):rs) = if i < n then x else go (i-n) rs
78   go i [] = error "miscalculation"
79   total = sum (M.elems m)
80   i = floor (r * fromIntegral total)
81
82 randData :: forall a. Data a => TyMap ConMap -> Ctxt ->  Rand StdGen a
83 randData tm ctxt = do
84     r <- getRandom 
85     let mc = selectRandom r cm
86         ctxt' i = extendCtxt x mc i ctxt
87     case mc of 
88         MyCon c -> fromConstrMI (\i -> randData tm (ctxt' i)) (repConstr (dataTypeOf x) c)
89         Str s -> return $ fromJust $ cast s
90   where
91     t = typeOf x
92     cm = tm M.! (t, shortenCtxt ctxt)
93     x = undefined :: a
94     
95 parseMode = defaultParseMode { fixities = Nothing }
96
97 main = do
98     args <- getArgs
99     modsParses <- mapM (parseFileWithMode parseMode) args
100     sequence_ [ hPutStrLn stderr $ prettyPrint l ++ ": " ++ e | ParseFailed l e <- modsParses]
101
102     let mods = [m | ParseOk m <- modsParses]
103
104     let cons = toTypeMap $ concatMap (getConOccs []) mods
105
106     m <- evalRandIO (randData cons []) :: IO Module
107
108     putStrLn (prettyPrint m)
109
110 data PairM m x = PM Int (m x)
111 unPairM (PM x m) = m
112
113 fromConstrMI :: forall m a. (Monad m, Data a)
114             => (forall d. Data d => Int -> m d)
115             -> Constr
116             -> m a
117 fromConstrMI f con = unPairM (gunfold k z con)
118  where
119   k :: forall b r. Data b => PairM m (b -> r) -> PairM m r
120   k (PM i c) = PM (i+1) (do { c' <- c ; b <- f i; return (c' b) })
121
122   z :: forall r. r -> PairM m r
123   z x = PM 0 (return x)
124