Fix eye function (#3)

Author: mdbrnowski

semantic_listener.py

@@ -170,7 +170,11 @@ def exitSpecialMatrixFunction(self, ctx: MyParser.SpecialMatrixFunctionContext):
                 type_dimentions.append(int(dim.getText()))
             else:
                 type_dimentions.append(None)
-        self.expr_type[ctx] = (Type.INT, *type_dimentions)
+        if ctx.getChild(0).symbol.type == MyParser.EYE:
+            assert len(type_dimentions) == 1
+            self.expr_type[ctx] = (Type.INT, type_dimentions[0], type_dimentions[0])
+        else:
+            self.expr_type[ctx] = (Type.INT, *type_dimentions)
 
     def exitVector(self, ctx: MyParser.VectorContext):
         elements = ctx.children[1::2]

tests/ast/input_3.txt

@@ -1,3 +1,3 @@
 a = 2 * (2 + 2) + 2 * 2;
 b = 3 * -3;
-C = 2 .+ 2 .* ones(3)';
+C = 2 .+ 2 .* eye(3)';

tests/ast/output_3.txt

@@ -22,5 +22,5 @@
 |  |  .*
 |  |  |  2
 |  |  |  transpose
-|  |  |  |  ones
+|  |  |  |  eye
 |  |  |  |  |  3

tests/semantic/input_special_matrix.txt

@@ -9,3 +9,5 @@ C = zeros(2, n);
 m = 3.14;
 
 D = eye(m);
+
+print eye(3) .+ ones(3, 3) .+ zeros(3, 3);