From d484fa1e5cf5b97ba1317a619eac48c272a3c44e Mon Sep 17 00:00:00 2001
From: pca006132 <john.lck40@gmail.com>
Date: Thu, 22 Jul 2021 11:49:00 +0800
Subject: [PATCH] added return type check

---
 nac3core/src/typecheck/type_inferencer/mod.rs  | 17 +++++++++++++++++
 nac3core/src/typecheck/type_inferencer/test.rs |  1 +
 nac3core/src/typecheck/typedef/mod.rs          |  2 +-
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/nac3core/src/typecheck/type_inferencer/mod.rs b/nac3core/src/typecheck/type_inferencer/mod.rs
index 624182b6d..45ef892d9 100644
--- a/nac3core/src/typecheck/type_inferencer/mod.rs
+++ b/nac3core/src/typecheck/type_inferencer/mod.rs
@@ -31,6 +31,7 @@ pub struct Inferencer<'a> {
     pub variable_mapping: HashMap<String, Type>,
     pub calls: &'a mut Vec<Rc<Call>>,
     pub primitives: &'a PrimitiveStore,
+    pub return_type: Option<Type>
 }
 
 struct NaiveFolder();
@@ -106,6 +107,20 @@ impl<'a> fold::Fold<()> for Inferencer<'a> {
                 }
             }
             ast::StmtKind::AnnAssign { .. } | ast::StmtKind::Expr { .. } => {}
+            ast::StmtKind::Return { value } => {
+                match (value, self.return_type) {
+                    (Some(v), Some(v1)) => {
+                        self.unifier.unify(v.custom.unwrap(), v1)?;
+                    }
+                    (Some(_), None) => {
+                        return Err("Unexpected return value".to_string());
+                    }
+                    (None, Some(_)) => {
+                        return Err("Expected return value".to_string());
+                    }
+                    (None, None) => {}
+                }
+            }
             _ => return Err("Unsupported statement type".to_string()),
         };
         Ok(stmt)
@@ -227,6 +242,7 @@ impl<'a> Inferencer<'a> {
             variable_mapping,
             calls: self.calls,
             primitives: self.primitives,
+            return_type: self.return_type
         };
         let fun = FunSignature {
             args: fn_args
@@ -275,6 +291,7 @@ impl<'a> Inferencer<'a> {
             variable_mapping,
             calls: self.calls,
             primitives: self.primitives,
+            return_type: self.return_type
         };
         let elt = new_context.fold_expr(elt)?;
         let generator = generators.pop().unwrap();
diff --git a/nac3core/src/typecheck/type_inferencer/test.rs b/nac3core/src/typecheck/type_inferencer/test.rs
index 3e85c2459..c85408f75 100644
--- a/nac3core/src/typecheck/type_inferencer/test.rs
+++ b/nac3core/src/typecheck/type_inferencer/test.rs
@@ -123,6 +123,7 @@ impl TestEnvironment {
             variable_mapping: Default::default(),
             calls: &mut self.calls,
             primitives: &mut self.primitives,
+            return_type: None
         }
     }
 }
diff --git a/nac3core/src/typecheck/typedef/mod.rs b/nac3core/src/typecheck/typedef/mod.rs
index 9a355c334..421f58ce7 100644
--- a/nac3core/src/typecheck/typedef/mod.rs
+++ b/nac3core/src/typecheck/typedef/mod.rs
@@ -177,7 +177,7 @@ impl Unifier {
             }
             TypeEnum::TObj { obj_id, params, .. } => {
                 let name = obj_to_name(*obj_id);
-                if params.len() > 0 {
+                if !params.is_empty() {
                     let mut params = params
                         .values()
                         .map(|v| self.stringify(*v, obj_to_name, var_to_name));