Skip to content

W
wasm-watermarker
Commit 8a26c932 authored by nagayama15's avatar nagayama15
Browse files
Options

Rewrite embedder using visitor

parent ff48a66b
Show whitespace changes
Inline Side-by-side
Showing with 193 additions and 199 deletions
src/kyut/watermarker/OperandSwappingWatermarker.cpp
View file @ 8a26c932
...@@ -4,6 +4,8 @@ ...@@ -4,6 +4,8 @@
#include <boost/optional.hpp> #include <boost/optional.hpp>
#include <wasm-traversal.h>
#include "../BitStreamWriter.hpp" #include "../BitStreamWriter.hpp"
#include "../CircularBitStreamReader.hpp" #include "../CircularBitStreamReader.hpp"
...@@ -317,117 +319,119 @@ namespace kyut::watermarker { ...@@ -317,117 +319,119 @@ namespace kyut::watermarker {
} }
// Watermark embedder // Watermark embedder
SideEffect embedExpression(wasm::Expression *expr, CircularBitStreamReader &stream); struct EmbeddingVisitor : wasm::OverriddenVisitor<EmbeddingVisitor, SideEffect> {
CircularBitStreamReader &stream;
explicit EmbeddingVisitor(CircularBitStreamReader &stream)
: stream(stream) {}
SideEffect embedExpressionList(const wasm::ExpressionList &exprs, CircularBitStreamReader &stream) { SideEffect visitExpressionList(const wasm::ExpressionList &exprs) {
auto effect = SideEffect::none; auto effect = SideEffect::none;
for (const auto expr : exprs) { for (const auto expr : exprs) {
effect = (std::max)(embedExpression(expr, stream), effect); effect = (std::max)(visit(expr), effect);
} }
return effect; return effect;
} }
SideEffect embedBlock(wasm::Block &expr, CircularBitStreamReader &stream) { SideEffect visitBlock(wasm::Block *expr) {
return embedExpressionList(expr.list, stream); return visitExpressionList(expr->list);
} }
SideEffect embedIf(wasm::If &expr, CircularBitStreamReader &stream) { SideEffect visitIf(wasm::If *expr) {
return (std::max)({ return (std::max)({
embedExpression(expr.condition, stream), visit(expr->condition),
embedExpression(expr.ifTrue, stream), visit(expr->ifTrue),
embedExpression(expr.ifFalse, stream), visit(expr->ifFalse),
}); });
} }
SideEffect embedLoop(wasm::Loop &expr, CircularBitStreamReader &stream) { SideEffect visitLoop(wasm::Loop *expr) {
return embedExpression(expr.body, stream); return visit(expr->body);
} }
SideEffect embedBreak(wasm::Break &expr, CircularBitStreamReader &stream) { SideEffect visitBreak(wasm::Break *expr) {
embedExpression(expr.value, stream); visit(expr->value);
embedExpression(expr.condition, stream); visit(expr->condition);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedSwitch(wasm::Switch &expr, CircularBitStreamReader &stream) { SideEffect visitSwitch(wasm::Switch *expr) {
return (std::max)(embedExpression(expr.condition, stream), embedExpression(expr.value, stream)); return (std::max)(visit(expr->condition), visit(expr->value));
} }
SideEffect embedCall(wasm::Call &expr, CircularBitStreamReader &stream) { SideEffect visitCall(wasm::Call *expr) {
embedExpressionList(expr.operands, stream); visitExpressionList(expr->operands);
// It is difficult to estimate the side effects of the function calls // It is difficult to estimate the side effects of the function calls
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedCallIndirect(wasm::CallIndirect &expr, CircularBitStreamReader &stream) { SideEffect visitCallIndirect(wasm::CallIndirect *expr) {
embedExpression(expr.target, stream); visit(expr->target);
embedExpressionList(expr.operands, stream); visitExpressionList(expr->operands);
// It is difficult to estimate the side effects of the function calls // It is difficult to estimate the side effects of the function calls
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedGetLocal([[maybe_unused]] wasm::GetLocal &expr, SideEffect visitGetLocal([[maybe_unused]] wasm::GetLocal *expr) {
[[maybe_unused]] CircularBitStreamReader &stream) {
return SideEffect::readOnly; return SideEffect::readOnly;
} }
SideEffect embedSetLocal(wasm::SetLocal &expr, CircularBitStreamReader &stream) { SideEffect visitSetLocal(wasm::SetLocal *expr) {
embedExpression(expr.value, stream); visit(expr->value);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedGetGlobal([[maybe_unused]] wasm::GetGlobal &expr, SideEffect visitGetGlobal([[maybe_unused]] wasm::GetGlobal *expr) {
[[maybe_unused]] CircularBitStreamReader &stream) {
return SideEffect::readOnly; return SideEffect::readOnly;
} }
SideEffect embedSetGlobal(wasm::SetGlobal &expr, CircularBitStreamReader &stream) { SideEffect visitSetGlobal(wasm::SetGlobal *expr) {
embedExpression(expr.value, stream); visit(expr->value);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedLoad(wasm::Load &expr, CircularBitStreamReader &stream) { SideEffect visitLoad(wasm::Load *expr) {
return (std::max)(embedExpression(expr.ptr, stream), SideEffect::readOnly); return (std::max)(visit(expr->ptr), SideEffect::readOnly);
} }
SideEffect embedStore(wasm::Store &expr, CircularBitStreamReader &stream) { SideEffect visitStore(wasm::Store *expr) {
embedExpression(expr.ptr, stream); visit(expr->ptr);
embedExpression(expr.value, stream); visit(expr->value);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedConst([[maybe_unused]] wasm::Const &expr, [[maybe_unused]] CircularBitStreamReader &stream) { SideEffect visitConst([[maybe_unused]] wasm::Const *expr) {
return SideEffect::none; return SideEffect::none;
} }
SideEffect embedUnary(wasm::Unary &expr, CircularBitStreamReader &stream) { SideEffect visitUnary(wasm::Unary *expr) {
return embedExpression(expr.value, stream); return visit(expr->value);
} }
SideEffect embedBinary(wasm::Binary &expr, CircularBitStreamReader &stream) { SideEffect visitBinary(wasm::Binary *expr) {
if (!isCommutative(expr.op)) { if (!isCommutative(expr->op)) {
// The operands of noncommutative instructions cannot be swapped // The operands of noncommutative instructions cannot be swapped
return (std::max)(embedExpression(expr.left, stream), embedExpression(expr.right, stream)); return (std::max)(visit(expr->left), visit(expr->right));
} }
if (!(*expr.left < *expr.right) && !(*expr.right < *expr.left)) { if (!(*expr->left < *expr->right) && !(*expr->right < *expr->left)) {
// If both sides are the same or cannot be ordered, skip embedding // If both sides are the same or cannot be ordered, skip visitding
return (std::max)(embedExpression(expr.left, stream), embedExpression(expr.right, stream)); return (std::max)(visit(expr->left), visit(expr->right));
} }
// Sort both of the operands // Sort both of the operands
auto [lo, hi] = std::minmax(expr.left, expr.right, [](auto a, auto b) { return *a < *b; }); auto [lo, hi] = std::minmax(expr->left, expr->right, [](auto a, auto b) { return *a < *b; });
const auto loEffect = embedExpression(lo, stream); const auto loEffect = visit(lo);
const auto hiEffect = embedExpression(hi, stream); const auto hiEffect = visit(hi);
if (static_cast<std::uint32_t>(loEffect) + static_cast<std::uint32_t>(hiEffect) >= 3) { if (static_cast<std::uint32_t>(loEffect) + static_cast<std::uint32_t>(hiEffect) >= 3) {
// The operands have side effect and cannot be swapped // The operands have side effect and cannot be swapped
...@@ -437,149 +441,138 @@ namespace kyut::watermarker { ...@@ -437,149 +441,138 @@ namespace kyut::watermarker {
// Embed watermarks by swapping operands // Embed watermarks by swapping operands
const bool bit = stream.readBit(); const bool bit = stream.readBit();
if (bit == (expr.left == lo)) { if (bit == (expr->left == lo)) {
swapOperands(expr); swapOperands(*expr);
} }
return (std::max)(loEffect, hiEffect); return (std::max)(loEffect, hiEffect);
} }
SideEffect embedSelect(wasm::Select &expr, CircularBitStreamReader &stream) { SideEffect visitSelect(wasm::Select *expr) {
return (std::max)({ return (std::max)({
embedExpression(expr.condition, stream), visit(expr->condition),
embedExpression(expr.ifTrue, stream), visit(expr->ifTrue),
embedExpression(expr.ifFalse, stream), visit(expr->ifFalse),
}); });
} }
SideEffect embedDrop(wasm::Drop &expr, CircularBitStreamReader &stream) { SideEffect visitDrop(wasm::Drop *expr) {
return embedExpression(expr.value, stream); return visit(expr->value);
} }
SideEffect embedReturn(wasm::Return &expr, CircularBitStreamReader &stream) { SideEffect visitReturn(wasm::Return *expr) {
embedExpression(expr.value, stream); visit(expr->value);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedHost(wasm::Host &expr, CircularBitStreamReader &stream) { SideEffect visitHost(wasm::Host *expr) {
embedExpressionList(expr.operands, stream); visitExpressionList(expr->operands);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedNop([[maybe_unused]] wasm::Nop &expr, [[maybe_unused]] CircularBitStreamReader &stream) { SideEffect visitNop([[maybe_unused]] wasm::Nop *exp) {
return SideEffect::none; return SideEffect::none;
} }
SideEffect embedUnreachable([[maybe_unused]] wasm::Unreachable &expr, SideEffect visitUnreachable([[maybe_unused]] wasm::Unreachable *expr) {
[[maybe_unused]] CircularBitStreamReader &stream) {
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedAtomicRMW(wasm::AtomicRMW &expr, CircularBitStreamReader &stream) { SideEffect visitAtomicRMW(wasm::AtomicRMW *expr) {
embedExpression(expr.ptr, stream); visit(expr->ptr);
embedExpression(expr.value, stream); visit(expr->value);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedAtomicCmpxchg(wasm::AtomicCmpxchg &expr, CircularBitStreamReader &stream) { SideEffect visitAtomicCmpxchg(wasm::AtomicCmpxchg *expr) {
embedExpression(expr.ptr, stream); visit(expr->ptr);
embedExpression(expr.expected, stream); visit(expr->expected);
embedExpression(expr.replacement, stream); visit(expr->replacement);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedAtomicWait(wasm::AtomicWait &expr, CircularBitStreamReader &stream) { SideEffect visitAtomicWait(wasm::AtomicWait *expr) {
embedExpression(expr.ptr, stream); visit(expr->ptr);
embedExpression(expr.expected, stream); visit(expr->expected);
embedExpression(expr.timeout, stream); visit(expr->timeout);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedAtomicNotify(wasm::AtomicNotify &expr, CircularBitStreamReader &stream) { SideEffect visitAtomicNotify(wasm::AtomicNotify *expr) {
embedExpression(expr.ptr, stream); visit(expr->ptr);
embedExpression(expr.notifyCount, stream); visit(expr->notifyCount);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedSIMDExtract(wasm::SIMDExtract &expr, CircularBitStreamReader &stream) { SideEffect visitSIMDExtract(wasm::SIMDExtract *expr) {
return embedExpression(expr.vec, stream); return visit(expr->vec);
} }
SideEffect embedSIMDReplace(wasm::SIMDReplace &expr, CircularBitStreamReader &stream) { SideEffect visitSIMDReplace(wasm::SIMDReplace *expr) {
return (std::max)(embedExpression(expr.vec, stream), embedExpression(expr.value, stream)); return (std::max)(visit(expr->vec), visit(expr->value));
} }
SideEffect embedSIMDShuffle(wasm::SIMDShuffle &expr, CircularBitStreamReader &stream) { SideEffect visitSIMDShuffle(wasm::SIMDShuffle *expr) {
return (std::max)(embedExpression(expr.left, stream), embedExpression(expr.right, stream)); return (std::max)(visit(expr->left), visit(expr->right));
} }
SideEffect embedSIMDBitselect(wasm::SIMDBitselect &expr, CircularBitStreamReader &stream) { SideEffect visitSIMDBitselect(wasm::SIMDBitselect *expr) {
return (std::max)({ return (std::max)({
embedExpression(expr.cond, stream), visit(expr->cond),
embedExpression(expr.left, stream), visit(expr->left),
embedExpression(expr.right, stream), visit(expr->right),
}); });
} }
SideEffect embedSIMDShift(wasm::SIMDShift &expr, CircularBitStreamReader &stream) { SideEffect visitSIMDShift(wasm::SIMDShift *expr) {
return (std::max)(embedExpression(expr.vec, stream), embedExpression(expr.shift, stream)); return (std::max)(visit(expr->vec), visit(expr->shift));
} }
SideEffect embedMemoryInit(wasm::MemoryInit &expr, CircularBitStreamReader &stream) { SideEffect visitMemoryInit(wasm::MemoryInit *expr) {
embedExpression(expr.dest, stream); visit(expr->dest);
embedExpression(expr.offset, stream); visit(expr->offset);
embedExpression(expr.size, stream); visit(expr->size);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedDataDrop([[maybe_unused]] wasm::DataDrop &expr, SideEffect visitDataDrop([[maybe_unused]] wasm::DataDrop *expr) {
[[maybe_unused]] CircularBitStreamReader &stream) {
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedMemoryCopy(wasm::MemoryCopy &expr, CircularBitStreamReader &stream) { SideEffect visitMemoryCopy(wasm::MemoryCopy *expr) {
embedExpression(expr.dest, stream); visit(expr->dest);
embedExpression(expr.source, stream); visit(expr->source);
embedExpression(expr.size, stream); visit(expr->size);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedMemoryFill(wasm::MemoryFill &expr, CircularBitStreamReader &stream) { SideEffect visitMemoryFill(wasm::MemoryFill *expr) {
embedExpression(expr.dest, stream); visit(expr->dest);
embedExpression(expr.value, stream); visit(expr->value);
embedExpression(expr.size, stream); visit(expr->size);
return SideEffect::write; return SideEffect::write;
} }
SideEffect embedExpression(wasm::Expression *expr, CircularBitStreamReader &stream) { SideEffect visit(wasm::Expression *expr) {
if (expr == nullptr) { if (expr == nullptr) {
return SideEffect::none; return SideEffect::none;
} }
return OverriddenVisitor::visit(expr);
switch (expr->_id) {
#define EXPR_TYPE(name) \
case ::wasm::Expression::name##Id: \
return embed##name(*expr->cast<::wasm::name>(), stream);
EXPR_TYPES()
#undef EXPR_TYPE
default:
WASM_UNREACHABLE();
}
} }
void embedFunction(wasm::Function &function, CircularBitStreamReader &stream) { void visitFunction(wasm::Function *function) {
embedExpression(function.body, stream); visit(function->body);
} }
};
} // namespace } // namespace
std::size_t embedOperandSwapping(wasm::Module &module, CircularBitStreamReader &stream) { std::size_t embedOperandSwapping(wasm::Module &module, CircularBitStreamReader &stream) {
...@@ -596,10 +589,11 @@ namespace kyut::watermarker { ...@@ -596,10 +589,11 @@ namespace kyut::watermarker {
std::begin(functions), std::end(functions), [](const auto a, const auto b) { return a->name < b->name; }); std::begin(functions), std::end(functions), [](const auto a, const auto b) { return a->name < b->name; });
// Embed watermarks // Embed watermarks
std::size_t posStart = stream.tell(); EmbeddingVisitor embedder{stream};
const auto posStart = stream.tell();
for (const auto f : functions) { for (const auto f : functions) {
embedFunction(*f, stream); embedder.visitFunction(f);
} }
return stream.tell() - posStart; return stream.tell() - posStart;
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment