From 00cf7eb0a1316ce0e330ed9a1e218fc60ae0b8cd Mon Sep 17 00:00:00 2001
From: Joel Kronqvist <joel.kronqvist@iki.fi>
Date: Sun, 3 Aug 2025 13:07:36 +0300
Subject: Implemented type_check, added method parts for SExp

'parts' is different from 'into_vec', because it doesn't require flat structure and always succeeds
---
 src/sexp/util.rs  | 12 +++++++++++
 src/type/check.rs | 64 +++++++++++++++++++++++++++++++++++++++++++++++++------
 src/type/mod.rs   |  2 +-
 3 files changed, 70 insertions(+), 8 deletions(-)

diff --git a/src/sexp/util.rs b/src/sexp/util.rs
index 70ac152..0818e6b 100644
--- a/src/sexp/util.rs
+++ b/src/sexp/util.rs
@@ -58,4 +58,16 @@ impl SExp {
             _ => Err("expected list, found atom".to_string()),
         }
     }
+
+    pub fn parts(self) -> Vec<SExp> {
+        match self {
+            Atom(x) => vec![Atom(x)],
+            SCons(a, b) if *b == Atom(Nil) => vec![*a],
+            SCons(a, b) => {
+                let mut res = vec![*a];
+                res.extend_from_slice(&(*b).parts());
+                res
+            },
+        }
+    }
 }
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))
     }
 
 }
diff --git a/src/type/mod.rs b/src/type/mod.rs
index 12d1132..5e7076a 100644
--- a/src/type/mod.rs
+++ b/src/type/mod.rs
@@ -8,7 +8,7 @@ pub mod check;
 use crate::sexp::SExp;
 
 
-#[derive(Debug,PartialEq)]
+#[derive(Clone,Debug,PartialEq)]
 pub enum Type {
 
     Integer,
-- 
cgit v1.2.3