-- A demonstration of lambda-lifting and -dropping

import List
import Test.QuickCheck
-- import Foreign      -- for debugging

data Term = Appl Term Term | Abstn Bool String Term | Loc Char Term 
          | Var String -- | Hole
            deriving (Eq, Show)

-- instance Show Term where
--     show (Var x) = " " ++ x
--     show (Appl m n ) = "(" ++ show m ++ " " ++ show n ++ ")"
--     show (Abstn True x m ) = "(%"++x++"."++ show m++")"
--     show (Abstn False x m ) = "(\\"++x++"."++ show m++")"
--     show (Loc l m) = "["++show m++"]" ++ [l]
-- --    show Hole = "[_]"

data Def = Def (String, Term)

instance Show Def where
    show (Def(f, t)) = f ++ " = " ++ show t

remove list x = [y | y <- list, x /= y]

freevars (Appl m n) = freevars m ++ freevars n
freevars (Abstn a x m) = remove (freevars m) x
freevars (Loc l m) = freevars m
freevars (Var x) = [x]
--freevars Hole = []

undup = nub

-- undup [] = []
-- undup (x:xs) = x : undup xs'
--     where xs' = remove xs x

lambdatower a vars m = lambdatower' a (reverse vars) m
    where lambdatower' a [] m = m
          lambdatower' a vars m = lambdatower' a xs (Abstn a x m)
              where x:xs = vars


appltower f [] = f
appltower f (arg:args) = appltower (Appl f arg) args 

close (Appl m n) = Appl (close m) (close n)
close (Abstn a x m) = appltower (lambdatower True fvs (Abstn a x m')) 
                      (map Var fvs)
    where fvs = undup $ freevars (Abstn a x m')
          m' = close' m
close (Loc l m) = appltower(Loc l (lambdatower True fvs m')) (map Var fvs)
    where fvs = undup $ freevars (Loc l m)
          m' = close' m
close m = m

close' (Abstn a x m) = Abstn a x (close' m)
close' m = close m

-- A monad that provides a gensym operator to manufacture new symbols.
data Gensym a = Gensym (Integer -> (a, Integer))

instance Monad Gensym where
    return x = Gensym (\y -> (x, y))
    (Gensym ma) >>= f = Gensym (\y -> let (a', y') = (ma y) in 
                                      let Gensym r = f a' in
                                      r y' )

gensym :: Gensym String
gensym = Gensym(\y -> ("f"++show y, y+1))

runGensym (Gensym f) = let (r, _) = f 0 in r

-- Formatting

times n x = take n $ repeat x

indent :: Int -> String -> String
indent n s = times n ' ' ++ s

indentedLine :: String -> String
indentedLine str = indent 4 (str ++ "\n")

printLifted term defs = 
    putStr $ show term ++ "\n" ++ "  where\n" ++
           concat (map indentedLine $ map show defs)

-- Lambda-Lifting

lift :: Char -> Term -> Gensym(Term, [Def])
lift l (Abstn a x m) = 
    do (m', defs) <- lift' l m
       name <- gensym
       return (Var name, Def(name, Loc l (Abstn a x m')):defs)
lift _ (Loc l m) = 
    do (m', defs) <- lift' l m
       name <- gensym
       return (Var name, Def(name, Loc l m'):defs)
lift l (Appl m n) = 
    do (m', mdefs) <- lift l m
       (n', ndefs) <- lift l n
       return (Appl m' n', mdefs ++ ndefs)
lift _ m = return (m, [])

-- lift' skips lambda towers and calls back to lift on any non-lambda
lift' l (Abstn a x m) = 
    do (m', defs) <- lift' l m
       return (Abstn a x  m', defs)
lift' l m = lift l m

lambdaLift t = runGensym $ lift undefined $ close t

subst name term (Var x) | x == name = term
                        | otherwise = (Var x)
subst name term (Appl m n) = Appl (subst name term m) (subst name term n)
subst name term (Abstn a x m) | name /= x = Abstn a x (subst name term m)
                              | otherwise = (Abstn a x m)
subst name term (Loc l m) = Loc l (subst name term m)
--subst name term Hole = Hole

inline [] m = m
inline (Def(name, body):defs) m = inline defs $ subst name body m

-- etaRedStep: a variation on eta-reduction
-- does just one reduction step, "etaReduce" is the trans. closure
--etaRedStep (Appl (Abstn True x m) (Var y)) = subst x (Var y) m
etaRedStep (Appl (Abstn True x m) (Var y)) | x == y = m
etaRedStep (Appl (Loc l (Abstn True x m)) (Var y)) = subst x (Var y) (Loc l m)
etaRedStep (Appl m n) = Appl (etaRedStep m) (etaRedStep n)
etaRedStep (Loc l m) = Loc l (etaRedStep m)
etaRedStep (Abstn a x m) = Abstn a x (etaRedStep m)
etaRedStep m = m

fixPoint f t = let t' = f t in 
               if t == t' then t 
               else fixPoint f t'

etaReduce = fixPoint etaRedStep

normalizeBrackets l' (Loc l m) | l == l' = normalizeBrackets l m
                                     | True    = Loc l (normalizeBrackets l m)
normalizeBrackets l (Appl m n) = Appl (normalizeBrackets l m)
                                            (normalizeBrackets l n)
normalizeBrackets l (Abstn a x m) = Abstn a x (normalizeBrackets l m)
normalizeBrackets l m = m

-- Tests!

unclosedDefs defs =  [(name, body) | Def(name, body) <- defs, 
                                     freevars body \\ names /= []]
    where names = [name | Def(name,_)<-defs]

testClosedDefs defs = unclosedDefs defs == []

testLiftDrop (s @ (Loc l m)) =
    let t = normalizeBrackets ' ' s in
    let (body, defs) = runGensym $ lift undefined $ close t in
    let t' = normalizeBrackets ' ' $ etaReduce $ inline defs body in
    (t == t', t, t')
testLiftDrop t = error ("Can only test terms with Loc at head")

liftDropInverse :: Term -> Bool
liftDropInverse (s @ (Loc l m)) =
    let t = normalizeBrackets ' ' s in
    let (body, defs) = runGensym $ lift undefined $ close t in
    let t' = normalizeBrackets ' ' $ etaReduce $ inline defs body in
    t == t'
liftDropInverse t = True

aTestTerm = (Loc 'c'
             (Abstn False "x" 
              (Appl(Abstn False "y" (Var "x"))
               (Abstn False "z" (Var "x")))))

containsLambda (Abstn a x m) = True
containsLambda (Appl m n) = containsLambda m || containsLambda n
containsLambda (Loc l m) = containsLambda m 
containsLambda m = False

containsLoc (Abstn a x m) = containsLoc m
containsLoc (Appl m n) = containsLambda m || containsLambda n
containsLoc (Loc l m) = True
containsLoc m = False

-- QuickCheck tests

variables = map (\x -> [x]) "abcdefghijklmnopqrstuvwxyz"
locations = "cs"

termGen :: Gen Term
termGen = oneof[fmap Var (elements variables),
                do { m <- termGen; n <- termGen; return (Appl m n) },
                do { m <- termGen; x <- elements variables;
                     return (Abstn False x m) },
                do { m <- termGen; l <- elements locations;
                     return (Loc l m)}
               ]

depthTermGen :: Int -> Gen Term
depthTermGen depth = 
    oneof $ [fmap Var (elements variables)] ++
          if (depth == 0) then [] else
              [do { m <- depthTermGen (depth-1); n <- depthTermGen (depth-1); 
                         return (Appl m n) },
               do { m <- depthTermGen (depth-1); x <- elements variables;
                         return (Abstn False x m) },
               do { m <- depthTermGen (depth-1); l <- elements locations;
                         return (Loc l m)}
              ]

sizedTermGen = sized depthTermGen

liftDropInverseProp = forAll sizedTermGen  $
                      \t -> classify (containsLambda t) "with lambda" $
                            classify (containsLoc t) "with loc" $
                            trivial (not(containsLambda t) && not(containsLoc t)) $
                            liftDropInverse t

liftedDefsClosed = forAll sizedTermGen $
                   \t -> let (t', defs) = lambdaLift t in
                         testClosedDefs defs

qcConf = Config 100 1 ((+3)) (\_ _ -> "") 

liftTest = check qcConf liftDropInverseProp