add a fastpath for recv[arg] when recv is untyped (#4760)

compiler/IREmitter/NameBasedIntrinsics.cc

@@ -632,6 +632,45 @@ class DefinedInstanceVar : public NameBasedIntrinsicMethod {
     }
 } DefinedInstanceVar;
 
+class SquareBrackets : public NameBasedIntrinsicMethod {
+public:
+    SquareBrackets() : NameBasedIntrinsicMethod{Intrinsics::HandleBlock::Unhandled} {};
+
+    // For cases where we don't have any type information, try to do something
+    // analogous to what the VM would do with the opt_aref instruction.
+    virtual llvm::Value *makeCall(MethodCallContext &mcctx) const override {
+        auto *send = mcctx.send;
+        if (send->args.size() != 1) {
+            return IREmitterHelpers::emitMethodCallViaRubyVM(mcctx);
+        }
+        // [] calls with blocks need to be handled differently.
+        if (mcctx.blk.has_value()) {
+            return IREmitterHelpers::emitMethodCallViaRubyVM(mcctx);
+        }
+        // If we had some kind of type information for the receiver, assume that we
+        // have already tested for a fast path earlier; this way we don't waste
+        // extra time doing another test that didn't work the first time.
+        if (!send->recv.type.isUntyped()) {
+            return IREmitterHelpers::emitMethodCallViaRubyVM(mcctx);
+        }
+
+        auto &cs = mcctx.cs;
+        auto &builder = mcctx.builder;
+        auto *cache = mcctx.getInlineCache();
+        auto *recv = mcctx.varGetRecv();
+        auto &args = mcctx.getStackArgs();
+        ENFORCE(args.size() == 1);
+
+        return builder.CreateCall(
+            cs.getFunction("sorbet_vm_aref"),
+            {Payload::getCFPForBlock(cs, builder, mcctx.irctx, mcctx.rubyBlockId), cache, recv, args[0]});
+    }
+
+    virtual InlinedVector<core::NameRef, 2> applicableMethods(CompilerState &cs) const override {
+        return {core::Names::squareBrackets()};
+    }
+} SquareBrackets;
+
 static const vector<CallCMethod> knownCMethods{
     {"<expand-splat>", "sorbet_expandSplatIntrinsic", NoReceiver, Intrinsics::HandleBlock::Unhandled,
      core::Symbols::Array()},
@@ -665,7 +704,7 @@ vector<const NameBasedIntrinsicMethod *> computeNameBasedIntrinsics() {
     vector<const NameBasedIntrinsicMethod *> ret{&DoNothingIntrinsic, &DefineClassIntrinsic,  &IdentityIntrinsic,
                                                  &CallWithBlock,      &ExceptionRetry,        &BuildHash,
                                                  &CallWithSplat,      &CallWithSplatAndBlock, &ShouldNeverSeeIntrinsic,
-                                                 &DefinedClassVar,    &DefinedInstanceVar};
+                                                 &DefinedClassVar,    &DefinedInstanceVar,    &SquareBrackets};
     for (auto &method : knownCMethods) {
         ret.emplace_back(&method);
     }

compiler/IREmitter/Payload/codegen-payload.c

@@ -193,6 +193,8 @@ SORBET_ALIVE(VALUE, sorbet_run_exception_handling,
 
 SORBET_ALIVE(VALUE, sorbet_rb_iterate,
              (VALUE(*body)(VALUE), VALUE data1, rb_block_call_func_t bl_proc, int minArgs, int maxArgs, VALUE data2));
+SORBET_ALIVE(VALUE, sorbet_vm_aref,
+             (rb_control_frame_t * cfp, struct FunctionInlineCache *cache, VALUE recv, VALUE arg));
 
 // The next several functions exist to convert Ruby definitions into LLVM IR, and
 // are always inlined as a consequence.

compiler/IREmitter/Payload/patches/vm_insnhelper.c

@@ -824,3 +824,17 @@ VALUE sorbet_rb_iterate(VALUE (*it_proc)(VALUE), VALUE data1, rb_block_call_func
     }
     return retval;
 }
+
+/* cf. vm_opt_aref */
+VALUE sorbet_vm_aref(rb_control_frame_t *reg_cfp, struct FunctionInlineCache *cache, VALUE recv, VALUE arg) {
+    if (!SPECIAL_CONST_P(recv)) {
+        if (RBASIC_CLASS(recv) == rb_cHash && BASIC_OP_UNREDEFINED_P(BOP_AREF, HASH_REDEFINED_OP_FLAG)) {
+            return rb_hash_aref(recv, arg);
+        }
+    }
+    VALUE *sp = reg_cfp->sp;
+    *(sp + 0) = recv;
+    *(sp + 1) = arg;
+    reg_cfp->sp += 2;
+    return sorbet_callFuncWithCache(cache, VM_BLOCK_HANDLER_NONE);
+}

test/testdata/compiler/intrinsics/hash_aref.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 # typed: true
 # compiled: true
+# run_filecheck: INITIAL
 
 extend T::Sig
 
@@ -15,13 +16,40 @@ def do_aref(x)
   x[:foo]
 end
 
+# Try and ensure that we inline for the hash case and don't call the vm-like
+# fastpath, since we have type information.
+
+# INITIAL-LABEL: @"func_Object#7do_aref"
+# INITIAL-NOT: call i64 @sorbet_vm_aref
+# INITIAL: call i64{{.*}}@sorbet_i_send
+# INITIAL-NOT: call i64 @sorbet_vm_aref
+# INITIAL: call i64 @sorbet_rb_hash_square_br
+# INITIAL-NOT: call i64 @sorbet_vm_aref
+# INITIAL{LITERAL}: }
+
+def do_aref_notype(x)
+  x[:foo]
+end
+
+# Make sure we call our vm-like fastpath for untyped args.
+
+# INITIAL-LABEL: @"func_Object#14do_aref_notype"
+# INITIAL-NOT: call i64 @sorbet_rb_hash_square_br
+# INITIAL: call i64 @sorbet_vm_aref
+# INITIAL-NOT: call i64 @sorbet_rb_hash_square_br
+# INITIAL{LITERAL}: }
+
 hash = {foo: 627}
 p hash
 
 3.times do
   y = do_aref(hash)
   p y
 
+  p do_aref_notype(hash)
+
   z = do_aref(NotHash.new)
   p z
+
+  p do_aref_notype(NotHash.new)
 end