52 changes: 34 additions & 18 deletions llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -145,34 +145,50 @@ class ObjectLinkingLayerJITLinkContext final : public JITLinkContext {
if (auto Err = MR.defineMaterializing(ExtraSymbolsToClaim))
return notifyFailed(std::move(Err));

if (const auto &InitSym = MR.getInitializerSymbol())
InternedResult[InitSym] = JITEvaluatedSymbol();

{

// Check that InternedResult matches up with MR.getSymbols().
// This guards against faulty transformations / compilers / object caches.

if (InternedResult.size() > MR.getSymbols().size()) {
SymbolNameVector ExtraSymbols;
for (auto &KV : InternedResult)
if (!MR.getSymbols().count(KV.first))
// First check that there aren't any missing symbols.
size_t NumMaterializationSideEffectsOnlySymbols = 0;
SymbolNameVector ExtraSymbols;
SymbolNameVector MissingSymbols;
for (auto &KV : MR.getSymbols()) {

// If this is a materialization-side-effects only symbol then bump
// the counter and make sure it's *not* defined, otherwise make
// sure that it is defined.
if (KV.second.hasMaterializationSideEffectsOnly()) {
++NumMaterializationSideEffectsOnlySymbols;
if (InternedResult.count(KV.first))
ExtraSymbols.push_back(KV.first);
ES.reportError(
make_error<UnexpectedSymbolDefinitions>(G.getName(),
std::move(ExtraSymbols)));
continue;
} else if (!InternedResult.count(KV.first))
MissingSymbols.push_back(KV.first);
}

// If there were missing symbols then report the error.
if (!MissingSymbols.empty()) {
ES.reportError(make_error<MissingSymbolDefinitions>(
G.getName(), std::move(MissingSymbols)));
MR.failMaterialization();
return;
}

SymbolNameVector MissingSymbols;
for (auto &KV : MR.getSymbols())
if (!InternedResult.count(KV.first))
MissingSymbols.push_back(KV.first);
// If there are more definitions than expected, add them to the
// ExtraSymbols vector.
if (InternedResult.size() >
MR.getSymbols().size() - NumMaterializationSideEffectsOnlySymbols) {
for (auto &KV : InternedResult)
if (!MR.getSymbols().count(KV.first))
ExtraSymbols.push_back(KV.first);
}

if (!MissingSymbols.empty()) {
ES.reportError(
make_error<MissingSymbolDefinitions>(G.getName(),
std::move(MissingSymbols)));
// If there were extra definitions then report the error.
if (!ExtraSymbols.empty()) {
ES.reportError(make_error<UnexpectedSymbolDefinitions>(
G.getName(), std::move(ExtraSymbols)));
MR.failMaterialization();
return;
}
Expand Down
3 changes: 0 additions & 3 deletions llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -258,9 +258,6 @@ Error RTDyldObjectLinkingLayer::onObjLoad(
Symbols.erase(KV.first);
}

if (const auto &InitSym = R.getInitializerSymbol())
Symbols[InitSym] = JITEvaluatedSymbol();

if (auto Err = R.notifyResolved(Symbols)) {
R.failMaterialization();
return Err;
Expand Down
56 changes: 56 additions & 0 deletions llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,62 @@ TEST_F(CoreAPIsStandardTest, EmptyLookup) {
EXPECT_TRUE(OnCompletionRun) << "OnCompletion was not run for empty query";
}

TEST_F(CoreAPIsStandardTest, ResolveUnrequestedSymbol) {
// Test that all symbols in a MaterializationUnit materialize corretly when
// only a subset of symbols is looked up.
// The aim here is to ensure that we're not relying on the query to set up
// state needed to materialize the unrequested symbols.

cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
SymbolFlagsMap({{Foo, FooSym.getFlags()}, {Bar, BarSym.getFlags()}}),
[this](MaterializationResponsibility R) {
cantFail(R.notifyResolved({{Foo, FooSym}, {Bar, BarSym}}));
cantFail(R.notifyEmitted());
})));

auto Result =
cantFail(ES.lookup(makeJITDylibSearchOrder(&JD), SymbolLookupSet({Foo})));
EXPECT_EQ(Result.size(), 1U) << "Unexpected number of results";
EXPECT_TRUE(Result.count(Foo)) << "Expected result for \"Foo\"";
}

TEST_F(CoreAPIsStandardTest, MaterializationSideEffctsOnlyTest) {
// Test that basic materialization-side-effects-only symbols work as expected:
// that they can be emitted without being resolved, that queries for them
// don't return until they're emitted, and that they don't appear in query
// results.

Optional<MaterializationResponsibility> FooR;
Optional<SymbolMap> Result;

cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
SymbolFlagsMap(
{{Foo, JITSymbolFlags::Exported |
JITSymbolFlags::MaterializationSideEffectsOnly}}),
[&](MaterializationResponsibility R) { FooR.emplace(std::move(R)); })));

ES.lookup(
LookupKind::Static, makeJITDylibSearchOrder(&JD),
SymbolLookupSet({Foo}, SymbolLookupFlags::WeaklyReferencedSymbol),
SymbolState::Ready,
[&](Expected<SymbolMap> LookupResult) {
if (LookupResult)
Result = std::move(*LookupResult);
else
ADD_FAILURE() << "Unexpected lookup error: "
<< toString(LookupResult.takeError());
},
NoDependenciesToRegister);

EXPECT_FALSE(Result) << "Lookup returned unexpectedly";
EXPECT_TRUE(FooR) << "Lookup failed to trigger materialization";
EXPECT_THAT_ERROR(FooR->notifyEmitted(), Succeeded())
<< "Emission of materialization-side-effects-only symbol failed";

EXPECT_TRUE(Result) << "Lookup failed to return";
EXPECT_TRUE(Result->empty()) << "Lookup result contained unexpected value";
}

TEST_F(CoreAPIsStandardTest, RemoveSymbolsTest) {
// Test that:
// (1) Missing symbols generate a SymbolsNotFound error.
Expand Down