1 files changed, 10 insertions, 6 deletions

diff --git a/src/type/check.rs b/src/type/check.rs
index 1712446..d815c85 100644
--- a/src/type/check.rs
+++ b/src/type/check.rs
@@ -112,11 +112,11 @@ impl SExp {
     ///
     /// let varlist = scons(var("a"), Nil);
     /// let typelist = scons(Ty(Integer), Nil);
-    /// let ret = Atom(Ty(Integer));
-    /// let body = scons(var("a"), Nil);
+    /// let ret = scons(Ty(Integer), Nil);
+    /// let body = scons(Add, scons(var("a"), scons(1, Nil)));
     /// assert_eq!(
     ///     scons(Fun, scons(varlist, scons(typelist, scons(ret, scons(body, Nil))))).type_check(),
-    ///     Ok(arr(List(vec![Integer]), Integer))
+    ///     Ok(arr(Integer, Integer))
     /// );
     /// ```
     /// Two-variable:
@@ -127,7 +127,7 @@ impl SExp {
     /// };
     ///
     /// let varlist = scons(var("a"), scons(var("b"), Nil));
-    /// let typelist = scons(Ty(List(vec![Integer])), Nil);
+    /// let typelist = scons(Ty(Integer), scons(Ty(Integer), Nil));
     /// let ret = Atom(Ty(Boolean));
     /// let body = scons(Eq, varlist.clone());
     /// assert_eq!(
@@ -202,7 +202,7 @@ impl SExp {
     }
 
 
-    fn infer_type(&self, mut ctx: HashMap<String, Type>) -> Result<Type, TypeError> {
+    pub fn infer_type(&self, mut ctx: HashMap<String, Type>) -> Result<Type, TypeError> {
 
         match self {
 
@@ -232,7 +232,7 @@ impl SExp {
 	    Atom(Print)  => Ok(arr(vt("_"), List(vec![]))),
 	    Atom(Ty(_))  => Ok(TypeLit),
 	    Atom(Arr)    => Ok(arr(List(vec![TypeLit, TypeLit]), TypeLit)),
-	    Atom(Fun)    => todo!(),
+	    Atom(Fun)    => Err(FunAsAtom),
 
             SCons(op, l) => {
 
@@ -249,6 +249,10 @@ impl SExp {
                     return Err(LetAsOperator(scons(op.clone(), l.clone())));
                 }
 
+		// Anonymous functions
+		if scons(op.clone(), l.clone()).is_fun() {
+		    return scons(op.clone(), l.clone()).get_fun_type(ctx);
+		}
 
                 // Normal operation
                 let opertype = (*op).infer_type(ctx.clone())?;