1 files changed, 31 insertions, 2 deletions
diff --git a/source/mod_student/typechecker.ml b/source/mod_student/typechecker.ml
index 109d22f..0d4d76d 100644
--- a/source/mod_student/typechecker.ml
+++ b/source/mod_student/typechecker.ml
@@ -211,7 +211,26 @@ module Make() = struct
         failwith "not yet implemented"
 
       | STMT_Return {values;loc} ->
-        failwith "not yet implemented"
+        begin
+          match TypingEnvironment.get_return env with
+          | None -> begin
+                      match values with
+                      | [] -> env,RT_Void
+                      | _ -> ErrorReporter.report_procedure_cannot_return_value ~loc
+                    end
+          | Some ret ->(
+            let expected = List.length ret
+            and actual   = List.length values in
+              if actual==0 then ErrorReporter.report_function_must_return_something ~loc
+              else
+              let rec aux rets values = match rets,values with
+                | ([],[]) -> env, RT_Void
+                | ((r::rs) , (v::vs)) -> let _=check_expression env r v in aux rs vs
+                | ([],(x::xs)) -> ErrorReporter.report_bad_number_of_return_values ~loc ~expected ~actual
+                | ((x::xs),[]) -> ErrorReporter.report_bad_number_of_return_values ~loc ~expected ~actual
+              in aux ret values)
+
+        end
 
       | STMT_VarDecl {var; init} ->
         begin
@@ -227,10 +246,20 @@ module Make() = struct
         end
 
       | STMT_While {cond; body; _} ->
+        check_expression env TP_Bool cond;
+        (* TODO *)
         failwith "not yet implemented"
 
     and check_statement_block env (STMTBlock {body; _}) =
-        failwith "not yet implemented"
+        let rec aux env = function
+          | []  -> env,RT_Unit
+          | [x] -> check_statement env x
+          | x::xs ->
+          begin match check_statement env x with
+                | (env,RT_Unit) -> aux env xs
+                | (env,RT_Void) -> ErrorReporter.report_other_error ~loc:(location_of_statement x) ~descr:"Return not at the end of block - unreachable code"
+          end
+        in env,(snd (aux env body))
 
     (* --------------------------------------------------- *)
     (* Top-level funkcje *)