/*
* Copyright 2025 WebAssembly Community Group participants
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
//
// Instruments branch hints and their targets, adding logging that allows us to
// see if the hints were valid or not. We turn
//
// @metadata.branch.hint B
// if (condition) {
// X
// } else {
// Y
// }
//
// into
//
// @metadata.branch.hint B
// ;; log the actual runtime result (condition), the prediction (B), and the
// ;; ID (123), and return that result.
// if (log(condition, B, 123)) {
// X
// } else {
// Y
// }
//
// Concretely, we emit calls to this logging function:
//
// (import "fuzzing-support" "log-branch"
// (func $log-branch (param i32 i32 i32) (result i32))
// )
//
// This can be used to verify that branch hints are accurate, by implementing
// the import like this for example:
//
// imports['fuzzing-support']['log-branch'] = (actual, prediction, id) => {
// // We only care about truthiness of the expected and actual values.
// expected = +!!expected;
// actual = +!!actual;
// // Throw if the hint said this branch would be taken, but it was not, or
// // vice versa.
// if (expected != actual) throw `Bad branch hint! (${id})`;
// return actual;
// };
//
// A pass to delete branch hints is also provided, which finds instrumentations
// and the IDs in those calls, and deletes branch hints that were listed. For
// example,
//
// --delete-branch-hints=10,20
//
// would do this transformation:
//
// @metadata.branch.hint A
// if (log(condition, A, 10)) { // 10 matches one of 10,20
// X
// }
// @metadata.branch.hint B
// if (log(condition, B, 99)) { // 99 does not match
// Y
// }
//
// =>
//
// // Used to be a branch hint here, but it was deleted.
// if (log(condition, A, 10)) {
// X
// }
// @metadata.branch.hint B // this one is unmodified.
// if (log(condition, B, 99)) {
// Y
// }
//
// A pass to undo the instrumentation is also provided, which does
//
// if (log(condition, A, 123)) {
// X
// }
//
// =>
//
// if (condition) {
// X
// }
//
#include "ir/effects.h"
#include "ir/names.h"
#include "ir/properties.h"
#include "ir/utils.h"
#include "pass.h"
#include "support/string.h"
#include "wasm-builder.h"
#include "wasm.h"
namespace wasm {
namespace {
// The module and base names of our import.
static const Name MODULE("fuzzing-support");
static const Name BASE("log-branch");
// Finds our import, if it exists.
Name getLogBranchImport(Module* module) {
for (auto& func : module->functions) {
if (func->module == MODULE && func->base == BASE) {
return func->name;
}
}
return Name();
}
// The branch id, which increments as we go.
int branchId = 1;
struct InstrumentBranchHints
: public WalkerPass<PostWalker<InstrumentBranchHints>> {
// The internal name of our import.
Name logBranch;
void visitIf(If* curr) { processCondition(curr); }
void visitBreak(Break* curr) {
if (curr->condition) {
processCondition(curr);
}
}
// TODO: BrOn, but the condition there is not an i32
template<typename T> void processCondition(T* curr) {
if (curr->condition->type == Type::unreachable) {
// This branch is not even reached.
return;
}
auto likely = getFunction()->codeAnnotations[curr].branchLikely;
if (!likely) {
return;
}
Builder builder(*getModule());
// Pick an ID for this branch.
int id = branchId++;
// Instrument the condition.
auto* idConst = builder.makeConst(Literal(int32_t(id)));
auto* guess = builder.makeConst(Literal(int32_t(*likely)));
curr->condition =
builder.makeCall(logBranch, {curr->condition, guess, idConst}, Type::i32);
}
void doWalkModule(Module* module) {
if (auto existing = getLogBranchImport(module)) {
// This file already has our import. We nop it out, as whatever the
// current code does may be dangerous (it may log incorrect hints).
auto* func = module->getFunction(existing);
Builder builder(*module);
if (func->getSig().results == Type::none) {
func->body = builder.makeNop();
} else {
func->body = builder.makeUnreachable();
}
func->module = func->base = Name();
func->type = func->type.with(Exact);
}
// Add our import.
auto* func = module->addFunction(Builder::makeFunction(
Names::getValidFunctionName(*module, BASE),
Type(Signature({Type::i32, Type::i32, Type::i32}, Type::i32),
NonNullable,
Inexact),
{}));
func->module = MODULE;
func->base = BASE;
logBranch = func->name;
// Walk normally, using logBranch as we go.
PostWalker<InstrumentBranchHints>::doWalkModule(module);
// Update ref.func type changes.
ReFinalize().run(getPassRunner(), module);
ReFinalize().walkModuleCode(module);
}
};
// Helper class that provides basic utilities for identifying and processing
// instrumentation from InstrumentBranchHints.
template<typename Sub>
struct InstrumentationProcessor : public WalkerPass<PostWalker<Sub>> {
using Super = WalkerPass<PostWalker<Sub>>;
// The internal name of our import.
Name logBranch;
Sub* self() { return static_cast<Sub*>(this); }
void visitIf(If* curr) { self()->processCondition(curr); }
void visitBreak(Break* curr) {
if (curr->condition) {
self()->processCondition(curr);
}
}
// TODO: BrOn, but the condition there is not an i32
void doWalkModule(Module* module) {
logBranch = getLogBranchImport(module);
if (!logBranch) {
Fatal()
<< "No branch hint logging import found. Was this code instrumented?";
}
Super::doWalkModule(module);
}
};
struct DeleteBranchHints : public InstrumentationProcessor<DeleteBranchHints> {
using Super = InstrumentationProcessor<DeleteBranchHints>;
// The set of IDs to delete.
std::unordered_set<Index> idsToDelete;
std::optional<uint32_t> getBranchID(Expression* condition,
const PassOptions& passOptions,
Module& wasm) {
auto* call =
Properties::getFallthrough(condition, getPassOptions(), *getModule())
->dynCast<Call>();
if (!call || call->target != logBranch || call->operands.size() != 3) {
return std::nullopt;
}
auto* c = call->operands[2]->dynCast<Const>();
if (!c || c->type != Type::i32) {
return std::nullopt;
}
return c->value.geti32();
}
template<typename T> void processCondition(T* curr) {
if (auto id = getBranchID(curr->condition, getPassOptions(), *getModule());
id && idsToDelete.contains(*id)) {
// Remove the branch hint.
getFunction()->codeAnnotations[curr].branchLikely = std::nullopt;
}
}
void doWalkModule(Module* module) {
auto arg = getArgument(
"delete-branch-hints",
"DeleteBranchHints usage: wasm-opt --delete-branch-hints=10,20,30");
for (auto& str : String::Split(arg, String::Split::NewLineOr(","))) {
idsToDelete.insert(std::stoi(str));
}
Super::doWalkModule(module);
}
};
struct DeInstrumentBranchHints
: public WalkerPass<PostWalker<DeInstrumentBranchHints>> {
// The internal name of our import.
Name logBranch;
void visitCall(Call* curr) {
if (curr->target == logBranch) {
// Replace the call with its first operand (the original condition).
replaceCurrent(curr->operands[0]);
}
}
void doWalkModule(Module* module) {
logBranch = getLogBranchImport(module);
if (!logBranch) {
Fatal()
<< "No branch hint logging import found. Was this code instrumented?";
}
// Mark the log-branch import as having no side effects - we are removing it
// entirely here, and its effect should not stop us when we compute effects.
module->getFunction(logBranch)->effects =
std::make_shared<EffectAnalyzer>(getPassOptions(), *module);
WalkerPass<PostWalker<DeInstrumentBranchHints>>::doWalkModule(module);
}
};
} // anonymous namespace
Pass* createInstrumentBranchHintsPass() { return new InstrumentBranchHints(); }
Pass* createDeleteBranchHintsPass() { return new DeleteBranchHints(); }
Pass* createDeInstrumentBranchHintsPass() {
return new DeInstrumentBranchHints();
}
} // namespace wasm