From 441af16f9a9a39ab311444620c0edb8bd98cfe79 Mon Sep 17 00:00:00 2001
From: Varun Gandhi <varun.gandhi@sourcegraph.com>
Date: Tue, 9 Jul 2024 17:55:54 +0800
Subject: [PATCH] fix: Propagate locations correctly for enums

---
 cfg/builder/builder_walk.cc         |  2 +-
 rewriter/TEnum.cc                   |  7 +++++--
 test/scip/testdata/enum.rb          | 10 +++++++++
 test/scip/testdata/enum.snapshot.rb | 32 +++++++++++++++++++++++++++++
 4 files changed, 48 insertions(+), 3 deletions(-)
 create mode 100644 test/scip/testdata/enum.rb
 create mode 100644 test/scip/testdata/enum.snapshot.rb

diff --git a/cfg/builder/builder_walk.cc b/cfg/builder/builder_walk.cc
index 7081f0711..1dd4e761e 100644
--- a/cfg/builder/builder_walk.cc
+++ b/cfg/builder/builder_walk.cc
@@ -415,7 +415,7 @@ BasicBlock *CFGBuilder::walk(CFGContext cctx, ast::ExpressionPtr &what, BasicBlo
             [&](ast::Assign &a) {
                 LocalRef lhs;
                 if (auto lhsIdent = ast::cast_tree<ast::ConstantLit>(a.lhs)) {
-                    lhs = global2Local(cctx, lhsIdent->symbol, a.loc);
+                    lhs = global2Local(cctx, lhsIdent->symbol, a.lhs.loc());
                 } else if (auto lhsLocal = ast::cast_tree<ast::Local>(a.lhs)) {
                     lhs = cctx.inWhat.enterLocal(lhsLocal->localVariable);
                 } else if (auto ident = ast::cast_tree<ast::UnresolvedIdent>(a.lhs)) {
diff --git a/rewriter/TEnum.cc b/rewriter/TEnum.cc
index c1cbf83ac..7317caa33 100644
--- a/rewriter/TEnum.cc
+++ b/rewriter/TEnum.cc
@@ -150,7 +150,10 @@ std::optional<ProcessStatResult> processStat(core::MutableContext ctx, ast::Clas
 
     auto statLocZero = stat.loc().copyWithZeroLength();
     auto name = ctx.state.enterNameConstant(ctx.state.freshNameUnique(core::UniqueNameKind::TEnum, lhs->cnst, 1));
-    auto classCnst = ast::MK::UnresolvedConstant(statLocZero, ast::MK::EmptyTree(), name);
+    // For some reason, Sorbet uses a zero-length range here,
+    // but it seems like we need this for scip-ruby?
+    // https://github.com/sorbet/sorbet/pull/7092
+    auto classCnst = ast::MK::UnresolvedConstant(lhs->loc, ast::MK::EmptyTree(), name);
     ast::ClassDef::ANCESTORS_store parent;
     parent.emplace_back(klass->name.deepCopy());
     ast::ClassDef::RHS_store classRhs;
@@ -243,7 +246,7 @@ void TEnum::run(core::MutableContext ctx, ast::ClassDef *klass) {
         auto serializeReturnTypeClass = core::cast_type_nonnull<core::ClassType>(serializeReturnType);
         ast::ExpressionPtr return_type_ast = ast::MK::Constant(klass->declLoc, serializeReturnTypeClass.symbol);
         auto sig = ast::MK::Sig0(klass->declLoc, std::move(return_type_ast));
-        auto method = ast::MK::SyntheticMethod0(klass->loc, klass->declLoc, klass->loc, core::Names::serialize(),
+        auto method = ast::MK::SyntheticMethod0(klass->loc, klass->declLoc, klass->name.loc(), core::Names::serialize(),
                                                 ast::MK::RaiseTypedUnimplemented(klass->declLoc));
         ast::Send::ARGS_store nargs;
         ast::Send::Flags flags;
diff --git a/test/scip/testdata/enum.rb b/test/scip/testdata/enum.rb
new file mode 100644
index 000000000..e725d37d9
--- /dev/null
+++ b/test/scip/testdata/enum.rb
@@ -0,0 +1,10 @@
+# typed: struct
+
+class X < T::Enum
+  enums do
+    A = new("A")
+    B = new
+  end
+
+  All = T.let([A, B], T::Array[X])
+end
diff --git a/test/scip/testdata/enum.snapshot.rb b/test/scip/testdata/enum.snapshot.rb
new file mode 100644
index 000000000..841b99454
--- /dev/null
+++ b/test/scip/testdata/enum.snapshot.rb
@@ -0,0 +1,32 @@
+ # typed: struct
+ 
+ class X < T::Enum
+#      ^ reference [..] X#
+#      ^ definition [..] X#
+#      ^ definition [..] X#serialize().
+#          ^ reference [..] T#
+#             ^^^^ reference [..] Module#public().
+#             ^^^^ reference [..] String#
+#             ^^^^ reference [..] T#Enum#
+   enums do
+     A = new("A")
+#    ^ definition local 2~#119448696
+#    ^ definition [..] X#A#
+#    ^ reference [..] X#A#
+#    ^ reference [..] X#A.
+#        ^^^ reference [..] Class#new().
+     B = new
+#    ^ definition local 5~#119448696
+#    ^ definition [..] X#B#
+#    ^ reference [..] X#B#
+#    ^ reference [..] X#B.
+#        ^^^ reference [..] Class#new().
+   end
+ 
+   All = T.let([A, B], T::Array[X])
+#  ^^^ definition [..] X#All.
+#               ^ reference [..] X#A.
+#                  ^ reference [..] X#B.
+#                         ^^^^^^^^ definition local 8~#119448696
+#                               ^ reference [..] X#
+ end