From 24fe4d183e21382a199367dfafa5cf3a6a8a2fb9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ibrahim=20Numanagic=CC=81?= <ibrahimpasa@gmail.com>
Date: Mon, 30 Jan 2023 15:22:35 -0800
Subject: [PATCH] Fix #183

---
 codon/parser/visitors/typecheck/infer.cpp | 38 +++++++++++++-------
 test/parser/types.codon                   | 43 +++++++++++++++++++++++
 2 files changed, 68 insertions(+), 13 deletions(-)

diff --git a/codon/parser/visitors/typecheck/infer.cpp b/codon/parser/visitors/typecheck/infer.cpp
index 90b528ea..8866d9db 100644
--- a/codon/parser/visitors/typecheck/infer.cpp
+++ b/codon/parser/visitors/typecheck/infer.cpp
@@ -387,6 +387,7 @@ StmtPtr TypecheckVisitor::prepareVTables() {
                                 N<IntExpr>(ctx->cache->classRealizationCnt + 1))));
   initAllVT.ast->suite = suite;
   auto typ = initAllVT.realizations.begin()->second->type;
+  LOG_REALIZE("[poly] {} : {}", typ, *suite);
   typ->ast = initAllVT.ast.get();
   auto fx = realizeFunc(typ.get(), true);
 
@@ -402,30 +403,37 @@ StmtPtr TypecheckVisitor::prepareVTables() {
   suite = N<SuiteStmt>();
   for (auto &[_, cls] : ctx->cache->classes) {
     for (auto &[r, real] : cls.realizations) {
+      size_t vtSz = 0;
+      for (auto &[base, vtable] : real->vtables) {
+        if (!vtable.ir)
+          vtSz += vtable.table.size();
+      }
+      auto var = initFn.ast->args[0].name;
+      // p.__setitem__(real.ID) = Ptr[cobj](real.vtables.size() + 2)
+      suite->stmts.push_back(N<ExprStmt>(N<CallExpr>(
+          N<DotExpr>(N<IdExpr>(var), "__setitem__"), N<IntExpr>(real->id),
+          N<CallExpr>(NT<InstantiateExpr>(NT<IdExpr>("Ptr"),
+                                          std::vector<ExprPtr>{NT<IdExpr>("cobj")}),
+                      N<IntExpr>(vtSz + 2)))));
+      // __internal__.class_set_typeinfo(p[real.ID], real.ID)
+      suite->stmts.push_back(N<ExprStmt>(
+          N<CallExpr>(N<IdExpr>("__internal__.class_set_typeinfo:0"),
+                      N<IndexExpr>(N<IdExpr>(var), N<IntExpr>(real->id)),
+                      N<IntExpr>(real->id))));
+      vtSz = 0;
       for (auto &[base, vtable] : real->vtables) {
         if (!vtable.ir) {
-          auto var = initFn.ast->args[0].name;
-          // p.__setitem__(real.ID) = Ptr[cobj](real.vtables.size() + 2)
-          suite->stmts.push_back(N<ExprStmt>(N<CallExpr>(
-              N<DotExpr>(N<IdExpr>(var), "__setitem__"), N<IntExpr>(real->id),
-              N<CallExpr>(NT<InstantiateExpr>(NT<IdExpr>("Ptr"),
-                                              std::vector<ExprPtr>{NT<IdExpr>("cobj")}),
-                          N<IntExpr>(vtable.table.size() + 2)))));
-          // __internal__.class_set_typeinfo(p[real.ID], real.ID)
-          suite->stmts.push_back(N<ExprStmt>(
-              N<CallExpr>(N<IdExpr>("__internal__.class_set_typeinfo:0"),
-                          N<IndexExpr>(N<IdExpr>(var), N<IntExpr>(real->id)),
-                          N<IntExpr>(real->id))));
           for (auto &[k, v] : vtable.table) {
             auto &[fn, id] = v;
             std::vector<ExprPtr> ids;
             for (auto &t : fn->getArgTypes())
               ids.push_back(NT<IdExpr>(t->realizedName()));
             // p[real.ID].__setitem__(f.ID, Function[<TYPE_F>](f).__raw__())
+            LOG_REALIZE("[poly] vtable[{}][{}] = {}", real->id, vtSz + id, fn);
             suite->stmts.push_back(N<ExprStmt>(N<CallExpr>(
                 N<DotExpr>(N<IndexExpr>(N<IdExpr>(var), N<IntExpr>(real->id)),
                            "__setitem__"),
-                N<IntExpr>(id),
+                N<IntExpr>(vtSz + id),
                 N<CallExpr>(N<DotExpr>(
                     N<CallExpr>(
                         NT<InstantiateExpr>(
@@ -438,12 +446,14 @@ StmtPtr TypecheckVisitor::prepareVTables() {
                         N<IdExpr>(fn->realizedName())),
                     "__raw__")))));
           }
+          vtSz += vtable.table.size();
         }
       }
     }
   }
   initFn.ast->suite = suite;
   typ = initFn.realizations.begin()->second->type;
+  LOG_REALIZE("[poly] {} : {}", typ, suite->toString(2));
   typ->ast = initFn.ast.get();
   realizeFunc(typ.get(), true);
 
@@ -469,6 +479,7 @@ StmtPtr TypecheckVisitor::prepareVTables() {
                 N<DotExpr>(N<IdExpr>(clsTyp->realizedName()), "__vtable_id__"))));
       }
 
+    LOG_REALIZE("[poly] {} : {}", t, *suite);
     initObjFns.ast->suite = suite;
     t->ast = initObjFns.ast.get();
     realizeFunc(t.get(), true);
@@ -502,6 +513,7 @@ StmtPtr TypecheckVisitor::prepareVTables() {
         N<DotExpr>(NT<InstantiateExpr>(
                        NT<IdExpr>(format("{}{}", TYPE_TUPLE, types.size())), types),
                    "__elemsize__"));
+    LOG_REALIZE("[poly] {} : {}", t, *suite);
     initDist.ast->suite = suite;
     t->ast = initDist.ast.get();
     realizeFunc(t.get(), true);
diff --git a/test/parser/types.codon b/test/parser/types.codon
index 8319646a..389b12a0 100644
--- a/test/parser/types.codon
+++ b/test/parser/types.codon
@@ -1775,3 +1775,46 @@ class Div(BinOp):
 expr : Expr = Mul(Const(3), Add(Const(10), Const(5)))
 print(expr.eval()) #: 45
 
+
+#%% polymorphism_4
+class A(object):
+    a: int
+    def __init__(self, a: int):
+        self.a = a
+
+    def test_a(self, n: int):
+        print("test_a:A", n)
+
+    def test(self, n: int):
+        print("test:A", n)
+
+    def test2(self, n: int):
+        print("test2:A", n)
+
+class B(A):
+    b: int
+    def __init__(self, a: int, b: int):
+        super().__init__(a)
+        self.b = b
+
+    def test(self, n: int):
+        print("test:B", n)
+
+    def test2(self, n: int):
+        print("test2:B", n)
+
+class C(B):
+    pass
+
+b = B(1, 2)
+b.test_a(1)
+b.test(1)
+#: test_a:A 1
+#: test:B 1
+
+a: A = b
+a.test(1)
+a.test2(2)
+#: test:B 1
+#: test2:B 2
+