1 files changed, 66 insertions, 8 deletions

diff --git a/src/type/check.rs b/src/type/check.rs
index 6e7e557..647ec4d 100644
--- a/src/type/check.rs
+++ b/src/type/check.rs
@@ -129,13 +129,16 @@ impl SExp {
             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
+            Atom(Quote) => Ok(arr(
+                               VarType("T".to_string()),
+                               VarType("T".to_string())
+                           )),
 
             SCons(op, l) => {
                 let opertype = (*op).infer_type(ctx.clone())?;
                 let argstype =  (*l).infer_list_type(ctx)?;
 
-                let opertype = if opertype.is_concrete() {
+                let opertype = if opertype.is_concrete().is_ok() {
                     opertype
                 } else {
                     opertype.infer_generics(&argstype)?
@@ -183,10 +186,66 @@ impl Type {
 
     /// Infers the type of generic 'VarType's using type of arguments.
     ///
-    /// In case there generic variables that can't be inferred,
+    /// In case there are generic variables that can't be inferred,
     /// returns an TypeError::UnboundVariable.
-    fn infer_generics(self, argtype: &Type) -> Result<Type, TypeError> {
-        todo!()
+    fn infer_generics(&self, argtype: &Type) -> Result<Type, TypeError> {
+        match self {
+            Arrow(from, to) => {
+
+                let generics = (*from).infer_generics_ctx(argtype, Vec::new())?;
+                let mut restype = (**to).clone();
+                for (name, ty) in generics {
+                    restype = restype.subst(&name, &ty);
+                }
+
+                match restype.is_concrete() {
+                    Ok(()) => Ok(arr(argtype.clone(), restype)),
+                    Err(unbound) => Err(UnboundGeneric(unbound)),
+                }
+            },
+            _ => Err(OtherError)
+        }
+    }
+
+    fn infer_generics_ctx(
+        &self,
+        argtype: &Type,
+        ctx: Vec<(String, Type)>
+    ) -> Result<Vec<(String, Type)>, TypeError> {
+        match (self, argtype) {
+
+            (a, b) if a == b => Ok(ctx),
+
+            (Arrow(a1, a2), Arrow(b1, b2)) => {
+                let mut r1 = a1.infer_generics_ctx(b1, ctx.clone())?;
+                let r2 = a2.infer_generics_ctx(b2, ctx.clone())?;
+                r1.extend_from_slice(&r2);
+                r1.extend_from_slice(&ctx);
+                Ok(r1)
+            },
+
+            (List(v1), List(v2)) => {
+                let mut res = ctx.clone();
+                for (t1, t2) in v1.into_iter().zip(v2.into_iter()) {
+                    let newctx = t1.infer_generics_ctx(t2, ctx.clone())?;
+                    res.extend_from_slice(&newctx);
+                }
+                Ok(res)
+            },
+
+            (VarType(name), ty) => {
+                let mut res = ctx.clone();
+                res.push((name.clone(), ty.clone()));
+                Ok(res)
+            },
+
+            (_a, _b) => Err(InvalidArgList {
+                arglist: Atom(Var("undefined".to_string())), // TODO: hacky as heck
+                expected: self.clone(),
+                found: argtype.clone(),
+            }),
+
+        }
     }
 
 }
@@ -194,8 +253,7 @@ impl Type {
 
 #[cfg(test)]
 mod tests {
-    use super::{*, Type::*, TypeError};
-    use crate::r#type::util::*;
+    use super::{*, TypeError};
 
     #[test]
     fn test_failing_infer_generics() {
@@ -242,7 +300,7 @@ mod tests {
             arr(List(vec![VarType("A".to_string()), VarType("B".to_string())]), VarType("B".to_string()))
                 .infer_generics(&List(vec![Integer, arr(Integer, Integer)])),
             Ok(arr(
-                    List(vec![List(vec![Integer, arr(Integer, Integer)])]),
+                    List(vec![Integer, arr(Integer, Integer)]),
                     arr(Integer, Integer)
                 ))
         );