Add return type to MethodSignature

3 files changed, 22 insertions, 9 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/MethodCallNode.java b/src/main/java/org/zwobble/couscous/ast/MethodCallNode.java
index c284d1f..c4e0321 100644
--- a/src/main/java/org/zwobble/couscous/ast/MethodCallNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/MethodCallNode.java
@@ -78,7 +78,10 @@ public class MethodCallNode implements ExpressionNode {
     }
 
     public MethodSignature signature() {
-        return new MethodSignature(methodName, eagerMap(arguments, argument -> argument.getType()));
+        return new MethodSignature(
+            methodName,
+            eagerMap(arguments, argument -> argument.getType()),
+            type);
     }
     
     @Override
diff --git a/src/main/java/org/zwobble/couscous/ast/MethodNode.java b/src/main/java/org/zwobble/couscous/ast/MethodNode.java
index e0b2432..dd970c5 100644
--- a/src/main/java/org/zwobble/couscous/ast/MethodNode.java
+++ b/src/main/java/org/zwobble/couscous/ast/MethodNode.java
@@ -140,7 +140,8 @@ public class MethodNode implements CallableNode {
     public MethodSignature signature() {
         return new MethodSignature(
             name,
-            eagerMap(arguments, argument -> argument.getType()));
+            eagerMap(arguments, argument -> argument.getType()),
+            returnType);
     }
 
     public MethodNode mapBody(Function<List<StatementNode>, List<StatementNode>> function) {
diff --git a/src/main/java/org/zwobble/couscous/ast/MethodSignature.java b/src/main/java/org/zwobble/couscous/ast/MethodSignature.java
index fb2dc64..d7bc725 100644
--- a/src/main/java/org/zwobble/couscous/ast/MethodSignature.java
+++ b/src/main/java/org/zwobble/couscous/ast/MethodSignature.java
@@ -1,15 +1,16 @@
 package org.zwobble.couscous.ast;
 
 import java.util.List;
-import java.util.Objects;
 
 public class MethodSignature {
-    private String name;
-    private List<TypeName> arguments;
+    private final String name;
+    private final List<TypeName> arguments;
+    private final TypeName returnType;
 
-    public MethodSignature(String name, List<TypeName> arguments) {
+    public MethodSignature(String name, List<TypeName> arguments, TypeName returnType) {
         this.name = name;
         this.arguments = arguments;
+        this.returnType = returnType;
     }
 
     public String getName() {
@@ -25,6 +26,7 @@ public class MethodSignature {
         return "MethodSignature(" +
             "name='" + name + '\'' +
             ", arguments=" + arguments +
+            ", returnType=" + returnType +
             ')';
     }
 
@@ -32,13 +34,20 @@ public class MethodSignature {
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
+
         MethodSignature that = (MethodSignature) o;
-        return Objects.equals(name, that.name) &&
-            Objects.equals(arguments, that.arguments);
+
+        if (!name.equals(that.name)) return false;
+        if (!arguments.equals(that.arguments)) return false;
+        return returnType.equals(that.returnType);
+
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(name, arguments);
+        int result = name.hashCode();
+        result = 31 * result + arguments.hashCode();
+        result = 31 * result + returnType.hashCode();
+        return result;
     }
 }