1 files changed, 26 insertions, 5 deletions
diff --git a/src/Haskell/Env.hs b/src/Haskell/Env.hs
index 21b2d4b..6b74221 100644
--- a/src/Haskell/Env.hs
+++ b/src/Haskell/Env.hs
@@ -1,6 +1,7 @@
 module Haskell.Env where
 
 import Control.Monad
+import Data.List
 import qualified Data.Map.Strict as Map
 import Haskell.AST
 
@@ -8,6 +9,10 @@ import Haskell.AST
 data Env = Env { eDefs :: Map.Map Name Expr }
   deriving (Show)
 
+instance Pretty Env where
+    pretty (Env defs) =
+        intercalate "\n" [n ++ " = " ++ pretty e | (n, e) <- Map.assocs defs]
+
 emptyEnv :: Env
 emptyEnv = Env Map.empty
 
@@ -18,12 +23,28 @@ addAST :: Env -> AST -> Either String Env
 addAST env (AST tops) = foldM addTop env tops
 
 addTop :: Env -> Toplevel -> Either String Env
-addTop env (TopDef (Def n ex)) =
-    Right $ env { eDefs = Map.insert n ex (eDefs env) }
+addTop env (TopDef def) = addDef env def
 addTop _ _ = Left "Only plain top-level definitions supported for the moment"
 
-forget :: Name -> Env -> Either String Env
-forget name env =
-    if Map.member name (eDefs env)
+addDef :: Env -> Def -> Either String Env
+addDef env (Def n ex) =
+    if envContains env n
+        then Left $ "Name '" ++ n ++ "' already present in environment"
+        else Right $ env { eDefs = Map.insert n ex (eDefs env) }
+
+envContains :: Env -> Name -> Bool
+envContains env name = Map.member name (eDefs env)
+
+forget :: Env -> Name -> Either String Env
+forget env name =
+    if envContains env name
         then Right $ env { eDefs = Map.delete name (eDefs env) }
         else Left $ "Name '" ++ name ++ "' not found in environment"
+
+reAdd :: Env -> Def -> Either String Env
+reAdd env def@(Def name _) = forget env name >>= flip addDef def
+
+get :: Env -> Name -> Either String Expr
+get env name = case Map.lookup name (eDefs env) of
+    Nothing -> Left "Name doesn't exist"
+    Just ex -> Right ex