module HW (
    fv,
    subst,
    normalstep,
    pickFresh,
    repeatedly,
    printnormal,
) where

import AST
import qualified Data.Set as Set
import Parse (parse)

-- | Return the free variables in an expression
fv :: Expr -> Set.Set String
fv (Var x) = Set.singleton x
fv (Lam x m) = x `Set.delete` (fv m)
fv (App m1 m2) = (fv m1) `Set.union` (fv m2)

{- | Substitute n for x in e, avoiding name capture
    subst n x e     e[x := n]
-}
subst :: Expr -> String -> Expr -> Expr
subst n x (Var e)
    | x == e = n
    | otherwise = Var e
subst n x (App m1 m2) = App (subst n x m1) (subst n x m2)
subst n x (Lam y m)
    | x == y = Lam y m
    | y `Set.notMember` (fv n) = Lam y (subst n x m)
    | otherwise = Lam y' (subst n x (subst (Var y') y m))
  where
    y' = pickFresh ((fv n) `Set.union` (fv m)) y

-- | Take a single step in normal order reduction or return Nothing
normalstep :: Expr -> Maybe Expr
-- beta
normalstep (App (Lam x m) n) = Just (subst n x m)
-- body
normalstep (Lam x m) = case normalstep m of
    Just m' -> Just (Lam x m')
    Nothing -> Nothing
-- arg
normalstep (App m n) | normalstep m == Nothing = case normalstep n of
    Just n' -> Just (App m n')
    Nothing -> Nothing
-- func
normalstep (App m n) = case normalstep m of
    Just m' -> Just (App m' n)
    Nothing -> Nothing
-- No further reductions
normalstep _ = Nothing

{- | Return a "fresh" name not already in the set.
 Tries x' then x'', etc.
-}
pickFresh :: Set.Set String -> String -> String
pickFresh s = pickFresh'
  where
    pickFresh' n | n `Set.notMember` s = n
    pickFresh' n = pickFresh' $ n ++ "'"

{- | Repeatedly apply a function to transform a value, returning the list
 of steps it took.  The result list starts with the given initial value
-}
repeatedly :: (a -> Maybe a) -> a -> [a]
repeatedly f = repeatedly'
  where
    repeatedly' x =
        x : case f x of
            Nothing -> []
            Just y -> repeatedly' y

-- | Print out the series of normal order reduction steps
printnormal :: String -> IO ()
printnormal = mapM_ print . repeatedly normalstep . parse