@wip replace environments serialized across compiler peers with stubs

rir/src/R/destubCloenv.cpp

@@ -0,0 +1,17 @@
+//
+// Created by Jakob Hain on 10/23/23.
+//
+
+#include "destubCloenv.h"
+#include "R/Symbols.h"
+#include "R/r.h"
+
+namespace rir {
+
+SEXP destubCloenv(SEXP closure) {
+    return CLOENV(closure) == symbol::closureEnvStub
+               ? R_GlobalEnv
+               : CLOENV(closure);
+}
+
+} // namespace rir

rir/src/R/destubCloenv.h

@@ -0,0 +1,13 @@
+//
+// Created by Jakob Hain on 10/23/23.
+//
+
+#pragma once
+
+#include "R/r_incl.h"
+
+namespace rir {
+
+SEXP destubCloenv(SEXP closure);
+
+} // namespace rir

rir/src/R/symbol_list.h

@@ -90,7 +90,8 @@
     V(remove, "remove")                                                        \
     V(rm, "rm")                                                                \
     V(Recall, "Recall")                                                        \
-    V(expandDotsTrigger, "\x02expandDotsTrigger\x03")
+    V(expandDotsTrigger, "\x02expandDotsTrigger\x03")                          \
+    V(closureEnvStub, "\x02closureEnvStub\x03")
 /*
  * The expandDotsTrigger symbol uses unprintable characters in hopes the users
  * won't create it from R (however, they still can, eg. `as.name("\x1a")`).

rir/src/compiler/compiler.cpp

@@ -1,13 +1,12 @@
 #include "compiler.h"
 #include "R/RList.h"
+#include "R/destubCloenv.h"
 #include "pir/continuation.h"
 #include "pir/pir_impl.h"
 #include "rir2pir/rir2pir.h"
 #include "runtime/TypeFeedback.h"
-#include "utils/Map.h"
 #include "utils/measuring.h"
 
-#include "compiler/analysis/query.h"
 #include "compiler/analysis/verifier.h"
 #include "compiler/opt/pass_definitions.h"
 #include "compiler/opt/pass_scheduler.h"
@@ -42,7 +41,7 @@ void Compiler::compileClosure(SEXP closure, const std::string& name,
     fun->clearDisabledAssumptions(assumptions);
     assumptions = tbl->combineContextWith(assumptions);
 
-    auto frame = RList(FRAME(CLOENV(closure)));
+    auto frame = RList(FRAME(destubCloenv(closure)));
 
     std::string closureName = name;
     if (name.compare("") == 0) {

rir/src/compiler/pir/closure.cpp

@@ -1,4 +1,5 @@
 #include "closure.h"
+#include "R/destubCloenv.h"
 #include "closure_version.h"
 #include "continuation.h"
 #include "env.h"
@@ -30,7 +31,7 @@ void Closure::invariant() const {
     // closure (since the closure is then created at runtime).
     assert(origin_ || env == Env::notClosed());
     assert(!origin_ || TYPEOF(origin_) == CLOSXP);
-    assert(env == Env::notClosed() || env->rho == CLOENV(origin_));
+    assert(env == Env::notClosed() || env->rho == destubCloenv(origin_));
     assert(!origin_ || formals_.original() == FORMALS(origin_));
 }
 

rir/src/compiler/pir/module.cpp

@@ -1,5 +1,6 @@
 #include "module.h"
 
+#include "R/destubCloenv.h"
 #include "compiler/parameter.h"
 #include "compilerClientServer/CompilerClient.h"
 #include "compilerClientServer/CompilerServer.h"
@@ -33,10 +34,10 @@ Closure* Module::getOrDeclareRirClosure(const std::string& name, SEXP closure,
     // For Identification we use the real env, but for optimization we only use
     // the real environment if this is not an inner function. When it is an
     // inner function, then the env is expected to change over time.
-    auto id = Idx(f, getEnv(CLOENV(closure)));
+    auto id = Idx(f, getEnv(destubCloenv(closure)));
     auto env = f->flags().contains(Function::InnerFunction)
                    ? Env::notClosed()
-                   : getEnv(CLOENV(closure));
+                   : getEnv(destubCloenv(closure));
     if (!closures.count(id))
         closures[id] = new Closure(name, closure, f, env, userContext);
     // If the compiler server is running sometimes this false.

rir/src/serializeHash/globals.cpp

@@ -35,6 +35,7 @@ void initGlobals() {
     cppId2Global_->emplace("R_DotsSymbol", R_DotsSymbol);
     cppId2Global_->emplace("R_NamesSymbol", R_NamesSymbol);
     cppId2Global_->emplace("expandDotsTrigger", symbol::expandDotsTrigger);
+    cppId2Global_->emplace("closureEnvStub", symbol::closureEnvStub);
 
     globals_ = new std::vector<SEXP>();
     globalsSet_ = new std::unordered_set<SEXP>();

rir/src/serializeHash/hash/UUIDPool.cpp

@@ -446,7 +446,7 @@ SEXP UUIDPool::readItem(const ByteBuffer& buf, bool useHashes) {
     }
 
     // Read regular data
-    return deserialize(buf, SerialOptions{useHashes, useHashes, false, false, SerialOptions::SourcePools()});
+    return deserialize(buf, SerialOptions{useHashes, useHashes, false, nullptr, SerialOptions::SourcePools()});
 }
 
 void UUIDPool::writeItem(SEXP sexp, __attribute__((unused)) bool isChild,
@@ -458,7 +458,7 @@ void UUIDPool::writeItem(SEXP sexp, __attribute__((unused)) bool isChild,
     }
 
     // Write regular data
-    serialize(sexp, buf, SerialOptions{useHashes, useHashes, false, false, SerialOptions::SourcePools()});
+    serialize(sexp, buf, SerialOptions{useHashes, useHashes, false, nullptr, SerialOptions::SourcePools()});
 }
 
 } // namespace rir

rir/src/serializeHash/serialize/serialize.cpp

@@ -1,10 +1,12 @@
 #include "serialize.h"
 #include "R/Printing.h"
 #include "R/Protect.h"
+#include "R/Symbols.h"
 #include "R/disableGc.h"
 #include "compiler/parameter.h"
 #include "compilerClientServer/CompilerServer.h"
 #include "runtime/PoolStub.h"
+#include "serializeHash/globals.h"
 #include "serializeHash/hash/UUIDPool.h"
 #include "serializeHash/hash/hashAst.h"
 #include "traceSerialize.h"
@@ -29,19 +31,19 @@ static const uint64_t dataBound = 0xfedcba9876543210;
 static const uint64_t intBound = 0xfedcba9876543211;
 #endif
 
-SerialOptions SerialOptions::DeepCopy{false, false, false, false, SerialOptions::SourcePools()};
+SerialOptions SerialOptions::DeepCopy{false, false, false, nullptr, SerialOptions::SourcePools()};
 
 SerialOptions SerialOptions::CompilerServer(bool intern) {
-    return SerialOptions{intern, intern, false, true, SerialOptions::SourcePools()};
+    return SerialOptions{intern, intern, false, nullptr, SerialOptions::SourcePools()};
 }
 
 SerialOptions SerialOptions::CompilerClient(bool intern, Function* function,
                                             SEXP decompiledClosure) {
-    return SerialOptions{intern, intern, false, true, SerialOptions::SourcePools(function, decompiledClosure)};
+    return SerialOptions{intern, intern, false, CLOENV(decompiledClosure), SerialOptions::SourcePools(function, decompiledClosure)};
 }
 
-SerialOptions SerialOptions::CompilerClientRetrieve{false, true, false, true, SerialOptions::SourcePools()};
-SerialOptions SerialOptions::SourceAndFeedback{false, true, true, true, SerialOptions::SourcePools()};
+SerialOptions SerialOptions::CompilerClientRetrieve{false, true, false, nullptr, SerialOptions::SourcePools()};
+SerialOptions SerialOptions::SourceAndFeedback{false, true, true, nullptr, SerialOptions::SourcePools()};
 
 unsigned pir::Parameter::RIR_SERIALIZE_CHAOS =
     getenv("RIR_SERIALIZE_CHAOS") ? strtol(getenv("RIR_SERIALIZE_CHAOS"), nullptr, 10) : 0;
@@ -103,28 +105,24 @@ SerialOptions SerialOptions::deserializeCompatible(AbstractDeserializer& deseria
     SerialOptions options;
     options.useHashes = deserializer.readBytesOf<bool>();
     options.onlySourceAndFeedback = deserializer.readBytesOf<bool>();
-    options.skipEnvLocks = deserializer.readBytesOf<bool>();
     return options;
 }
 
 void SerialOptions::serializeCompatible(AbstractSerializer& serializer) const {
     serializer.writeBytesOf(useHashes);
     serializer.writeBytesOf(onlySourceAndFeedback);
-    serializer.writeBytesOf(skipEnvLocks);
 }
 
 bool SerialOptions::areCompatibleWith(const rir::SerialOptions& other) const {
     return useHashes == other.useHashes &&
-           onlySourceAndFeedback == other.onlySourceAndFeedback &&
-           skipEnvLocks == other.skipEnvLocks;
+           onlySourceAndFeedback == other.onlySourceAndFeedback;
 }
 
 bool SerialOptions::willReadOrWrite(const SerialFlags& flags) const {
     return
         (!onlySourceAndFeedback ||
          flags.contains(SerialFlag::InSource) ||
-         flags.contains(SerialFlag::InFeedback)) &&
-        (!skipEnvLocks || flags.contains(SerialFlag::NotEnvLock));
+         flags.contains(SerialFlag::InFeedback));
 }
 
 bool Serializer::willWrite(const rir::SerialFlags& flags) const {
@@ -170,6 +168,14 @@ void Serializer::write(SEXP s, const SerialFlags& flags) {
     // If this is a stubbed pool entry, serialize the stub instead
     if (options.sourcePools.isEntry(s)) {
         s = options.sourcePools.stub(s);
+    } else if (s == options.closureEnvAndIfSetWeTryToSerializeLocalEnvsAsStubs) {
+        s = symbol::closureEnvStub;
+    } else if (options.closureEnvAndIfSetWeTryToSerializeLocalEnvsAsStubs &&
+               TYPEOF(s) == ENVSXP && !globalsSet.count(s) &&
+               !R_IsPackageEnv(s) && !R_IsNamespaceEnv(s)) {
+        std::cerr << "WARNING: pointerStubLocalEnvs isn't implemented, and "
+                  << "we're serializing a local env: " << Print::dumpSexp(s)
+                  << std::endl;
     }
 
 #if DEBUG_SERIALIZE_CONSISTENCY
@@ -311,8 +317,11 @@ SEXP Deserializer::read(const SerialFlags& flags) {
            "serialize/deserialize sexp type mismatch");
 #endif
 
-    // If this is a stubbed pool entry, deserialize the stub instead
-    if (options.sourcePools.isStub(result)) {
+    // If this is a stub, deserialize the stubbed value instead
+    if (options.closureEnvAndIfSetWeTryToSerializeLocalEnvsAsStubs &&
+        result == symbol::closureEnvStub) {
+        result = options.closureEnvAndIfSetWeTryToSerializeLocalEnvsAsStubs;
+    } else if (options.sourcePools.isStub(result)) {
         result = options.sourcePools.entry(result);
     }
 

rir/src/serializeHash/serialize/serialize.h

@@ -43,8 +43,9 @@ struct SerialOptions {
     bool useHashesForRecordedCalls;
     /// Whether to only serialize source and feedback (no optimized code).
     bool onlySourceAndFeedback;
-    /// Whether to skip serializing environment locks
-    bool skipEnvLocks;
+    /// If set, will serialize this as a closure environment stub, and warn when
+    /// other local environments are serialized
+    SEXP closureEnvAndIfSetWeTryToSerializeLocalEnvsAsStubs;
     /// If nonempty, we serialize the corresponding SEXPs with stubs from these
     /// pools
     SourcePools sourcePools;

rir/src/serializeHash/serialize/serializeR.cpp

@@ -252,7 +252,7 @@ static void rStreamInBytes(R_inpstream_t stream, void* data, int length) {
 }
 
 static SerialOptions* newRSerialOptions(bool useHashes) {
-    return new SerialOptions{useHashes, useHashes, false, false, SerialOptions::SourcePools()};
+    return new SerialOptions{useHashes, useHashes, false, nullptr, SerialOptions::SourcePools()};
 }
 
 void serializeR(SEXP sexp, ByteBuffer& buffer, bool useHashes) {