From 797440387c1ec590914ccb45c1350dce2b549a92 Mon Sep 17 00:00:00 2001
From: bd-912 <bdunahu@colostate.edu>
Date: Fri, 26 Apr 2024 20:57:37 -0600
Subject: HeatVisitor.MethodDeclaraction checks correct return type

---
 heat/HeatVisitor.java | 17 ++++++++++++++---
 heat/TypeBundle.java  |  4 ++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/heat/HeatVisitor.java b/heat/HeatVisitor.java
index 8d29813..f099a43 100644
--- a/heat/HeatVisitor.java
+++ b/heat/HeatVisitor.java
@@ -240,11 +240,19 @@ public class HeatVisitor extends GJDepthFirst<TypeBundle,ArrayList<String>> {
                                          id));
         this.symt.setActive(TypeEnum.method, symt.getMethod(id));
         ///////////////////////////////////////////////////////////////
-        // n.f1.accept(this, argu);
+        TypeBundle tb = n.f1.accept(this, argu);
         // n.f4.accept(this, para);
         n.f7.accept(this, argu);
         n.f8.accept(this, argu);
-        n.f10.accept(this, argu);
+        _ret = n.f10.accept(this, argu);
+
+        // NOT
+        // if it is a class and the classes are equal
+        // or it is not a class and the types are equal
+        if (!((tb.getInstance() != null && tb.getInstance().equals(_ret.getInstance())) ||
+              (tb.getInstance() == null && tb.getType() == _ret.getType())))
+            throw new TypecheckException(String.format("%s returns the wrong type!",
+                                                       id));
         ///////////////////////////////////////////////////////////////
         this.symt.removeActive(TypeEnum.method);
         MinimalLogger.info(String.format("<- %s (%s)",
@@ -319,7 +327,10 @@ public class HeatVisitor extends GJDepthFirst<TypeBundle,ArrayList<String>> {
         MinimalLogger.info(String.format("-> %s",
                                          n.getClass().getSimpleName()));
         ///////////////////////////////////////////////////////////////
-        _ret = n.f0.accept(this, argu);
+        if (n.f0.which == 3)
+            _ret = new TypeBundle(TypeEnum.classname, this.symt.getClass(((Identifier) n.f0.choice).f0.tokenImage));
+        else
+            _ret = n.f0.accept(this, argu);
         ///////////////////////////////////////////////////////////////
         MinimalLogger.info(String.format("<- %s with %s",
                                          n.getClass().getSimpleName(),
diff --git a/heat/TypeBundle.java b/heat/TypeBundle.java
index d8bfb4f..ac00e71 100644
--- a/heat/TypeBundle.java
+++ b/heat/TypeBundle.java
@@ -53,4 +53,8 @@ class TypeBundle {
         return this.type;
     }
 
+    public ClassInstance getInstance() {
+        return this.instance;
+    }
+
 }
-- 
cgit v1.2.3