1 files changed, 57 insertions, 7 deletions
diff --git a/src/type/check.rs b/src/type/check.rs
index 3f3b864..8d2a35d 100644
--- a/src/type/check.rs
+++ b/src/type/check.rs
@@ -1,6 +1,6 @@
 
-use crate::r#type::{Type, TypeError};
-use crate::sexp::SExp;
+use crate::r#type::{Type, TypeError, Type::*, TypeError::*, util::*};
+use crate::sexp::{SExp, SExp::*, SLeaf::*};
 use std::collections::HashMap;
 
 impl SExp {
@@ -65,10 +65,12 @@ impl SExp {
     ///     sexp::{SExp::*, SLeaf::*, util::*},
     /// };
     ///
-    /// match scons(Mul, scons(1, Nil)).type_check() {
+    /// let err = scons(Mul, scons(1, Nil)).type_check();
+    /// match err {
     ///     Err(InvalidArgList { .. }) => (),
     ///     _ => panic!(
-    ///         "passing only 1 argument to '*' should result in InvalidArgList error"
+    ///         "passing only 1 argument to '*' should result in InvalidArgList error, found '{:?}'",
+    ///         err
     ///     ),
     /// };
     ///
@@ -110,12 +112,60 @@ impl SExp {
     /// };
     /// ```
     pub fn type_check(&self) -> Result<Type, TypeError> {
-        self.type_check_ctx(HashMap::new())
+        self.infer_type(HashMap::new())
     }
 
 
-    fn type_check_ctx(&self, ctx: HashMap<String, Type>) -> Result<Type, TypeError> {
-        todo!();
+    fn infer_type(&self, ctx: HashMap<String, Type>) -> Result<Type, TypeError> {
+
+        match self {
+
+            Atom(Int(_)) => Ok(Integer),
+            Atom(Var(name)) => ctx.get(name)
+                                  .ok_or(UndefinedVariable(name.to_string()))
+                                  .cloned(),
+            Atom(Add) => Ok(arr(List(vec![Integer, Integer]), Integer)), // TODO varlen
+            Atom(Mul) => Ok(arr(List(vec![Integer, Integer]), Integer)), // TODO varlen
+            Atom(Sub) => Ok(arr(List(vec![Integer, Integer]), Integer)),
+            Atom(Div) => Ok(arr(List(vec![Integer, Integer]), Integer)),
+            Atom(Nil)   => Ok(List(vec![])),
+            Atom(Quote) => Ok(List(vec![])), // TODO make it all identity functions
+
+            SCons(op, l) if **op == Atom(Quote) => (*l).infer_list_type(ctx),
+
+            SCons(op, l) => {
+                let opertype = (*op).infer_type(ctx.clone())?;
+                let argstype =  (*l).infer_list_type(ctx)?;
+                match (opertype, argstype) {
+                    (Arrow(a, b), c) => {
+                        if *a != c {
+                            Err(InvalidArgList {
+                                    arglist:  (**l).clone(),
+                                    expected: *a,
+                                    found:    c,
+                            })
+                        } else {
+                            Ok(*b)
+                        }
+                    },
+                    (t, _) => Err(InvalidOperator {
+                        operator: *op.clone(),
+                        expected: arr(UndefinedType, UndefinedType),
+                        found:    t,
+                    }),
+                }
+            },
+
+        }
+
+    }
+
+    fn infer_list_type(&self, ctx: HashMap<String, Type>) -> Result<Type, TypeError> {
+        let mut res = vec![];
+        for exp in self.clone().parts() {
+            res.push(exp.infer_type(ctx.clone())?);
+        }
+        Ok(List(res))
     }
 
 }