Skip to content

W
wasm-watermarker
Commit 1e3f6211 authored by nagayama15's avatar nagayama15
Browse files
Options

Refactor watermark embedder and extractor

parent 9962536d
Hide whitespace changes
Inline Side-by-side
Showing with 27 additions and 268 deletions
src/kyut/watermarker/OperandSwappingWatermarker.cpp
View file @ 1e3f6211
...@@ -318,12 +318,13 @@ namespace kyut::watermarker { ...@@ -318,12 +318,13 @@ namespace kyut::watermarker {
return false; return false;
} }
// Watermark embedder // Watermarking visitor
struct EmbeddingVisitor : wasm::OverriddenVisitor<EmbeddingVisitor, SideEffect> { template <typename F>
CircularBitStreamReader &stream; struct OperandSwappingVisitor : wasm::OverriddenVisitor<OperandSwappingVisitor<F>, SideEffect> {
F f;
explicit EmbeddingVisitor(CircularBitStreamReader &stream) explicit OperandSwappingVisitor(F f)
: stream(stream) {} : f(std::move(f)) {}
SideEffect visitExpressionList(const wasm::ExpressionList &exprs) { SideEffect visitExpressionList(const wasm::ExpressionList &exprs) {
auto effect = SideEffect::none; auto effect = SideEffect::none;
...@@ -438,12 +439,7 @@ namespace kyut::watermarker { ...@@ -438,12 +439,7 @@ namespace kyut::watermarker {
return (std::max)(loEffect, hiEffect); return (std::max)(loEffect, hiEffect);
} }
// Embed watermarks by swapping operands f(*expr);
const bool bit = stream.readBit();
if (bit == (expr->left == lo)) {
swapOperands(*expr);
}
return (std::max)(loEffect, hiEffect); return (std::max)(loEffect, hiEffect);
} }
...@@ -566,259 +562,7 @@ namespace kyut::watermarker { ...@@ -566,259 +562,7 @@ namespace kyut::watermarker {
if (expr == nullptr) { if (expr == nullptr) {
return SideEffect::none; return SideEffect::none;
} }
return OverriddenVisitor::visit(expr); return static_cast<wasm::OverriddenVisitor<OperandSwappingVisitor<F>, SideEffect> *>(this)->visit(expr);
}
void visitFunction(wasm::Function *function) {
visit(function->body);
}
};
// Watermark extractor
struct ExtractingVisitor : wasm::OverriddenVisitor<ExtractingVisitor, SideEffect> {
BitStreamWriter &stream;
explicit ExtractingVisitor(BitStreamWriter &stream)
: stream(stream) {}
SideEffect visitExpressionList(const wasm::ExpressionList &exprs) {
auto effect = SideEffect::none;
for (const auto expr : exprs) {
effect = (std::max)(visit(expr), effect);
}
return effect;
}
SideEffect visitBlock(wasm::Block *expr) {
return visitExpressionList(expr->list);
}
SideEffect visitIf(wasm::If *expr) {
return (std::max)({
visit(expr->condition),
visit(expr->ifTrue),
visit(expr->ifFalse),
});
}
SideEffect visitLoop(wasm::Loop *expr) {
return visit(expr->body);
}
SideEffect visitBreak(wasm::Break *expr) {
visit(expr->value);
visit(expr->condition);
return SideEffect::write;
}
SideEffect visitSwitch(wasm::Switch *expr) {
return (std::max)(visit(expr->condition), visit(expr->value));
}
SideEffect visitCall(wasm::Call *expr) {
visitExpressionList(expr->operands);
// It is difficult to estimate the side effects of the function calls
return SideEffect::write;
}
SideEffect visitCallIndirect(wasm::CallIndirect *expr) {
visit(expr->target);
visitExpressionList(expr->operands);
// It is difficult to estimate the side effects of the function calls
return SideEffect::write;
}
SideEffect visitGetLocal([[maybe_unused]] wasm::GetLocal *expr) {
return SideEffect::readOnly;
}
SideEffect visitSetLocal(wasm::SetLocal *expr) {
visit(expr->value);
return SideEffect::write;
}
SideEffect visitGetGlobal([[maybe_unused]] wasm::GetGlobal *expr) {
return SideEffect::readOnly;
}
SideEffect visitSetGlobal(wasm::SetGlobal *expr) {
visit(expr->value);
return SideEffect::write;
}
SideEffect visitLoad(wasm::Load *expr) {
return (std::max)(visit(expr->ptr), SideEffect::readOnly);
}
SideEffect visitStore(wasm::Store *expr) {
visit(expr->ptr);
visit(expr->value);
return SideEffect::write;
}
SideEffect visitConst([[maybe_unused]] wasm::Const *expr) {
return SideEffect::none;
}
SideEffect visitUnary(wasm::Unary *expr) {
return visit(expr->value);
}
SideEffect visitBinary(wasm::Binary *expr) {
if (!isCommutative(expr->op)) {
// The operands of noncommutative instructions cannot be swapped
return (std::max)(visit(expr->left), visit(expr->right));
}
if (!(*expr->left < *expr->right) && !(*expr->right < *expr->left)) {
// If both sides are the same or cannot be ordered, skip visitding
return (std::max)(visit(expr->left), visit(expr->right));
}
// Sort both of the operands
auto [lo, hi] = std::minmax(expr->left, expr->right, [](auto a, auto b) { return *a < *b; });
const auto loEffect = visit(lo);
const auto hiEffect = visit(hi);
if (static_cast<std::uint32_t>(loEffect) + static_cast<std::uint32_t>(hiEffect) >= 3) {
// The operands have side effect and cannot be swapped
return (std::max)(loEffect, hiEffect);
}
// Extract watermarks from operands orders
stream.writeBit(lo == expr->right);
return (std::max)(loEffect, hiEffect);
}
SideEffect visitSelect(wasm::Select *expr) {
return (std::max)({
visit(expr->condition),
visit(expr->ifTrue),
visit(expr->ifFalse),
});
}
SideEffect visitDrop(wasm::Drop *expr) {
return visit(expr->value);
}
SideEffect visitReturn(wasm::Return *expr) {
visit(expr->value);
return SideEffect::write;
}
SideEffect visitHost(wasm::Host *expr) {
visitExpressionList(expr->operands);
return SideEffect::write;
}
SideEffect visitNop([[maybe_unused]] wasm::Nop *exp) {
return SideEffect::none;
}
SideEffect visitUnreachable([[maybe_unused]] wasm::Unreachable *expr) {
return SideEffect::write;
}
SideEffect visitAtomicRMW(wasm::AtomicRMW *expr) {
visit(expr->ptr);
visit(expr->value);
return SideEffect::write;
}
SideEffect visitAtomicCmpxchg(wasm::AtomicCmpxchg *expr) {
visit(expr->ptr);
visit(expr->expected);
visit(expr->replacement);
return SideEffect::write;
}
SideEffect visitAtomicWait(wasm::AtomicWait *expr) {
visit(expr->ptr);
visit(expr->expected);
visit(expr->timeout);
return SideEffect::write;
}
SideEffect visitAtomicNotify(wasm::AtomicNotify *expr) {
visit(expr->ptr);
visit(expr->notifyCount);
return SideEffect::write;
}
SideEffect visitSIMDExtract(wasm::SIMDExtract *expr) {
return visit(expr->vec);
}
SideEffect visitSIMDReplace(wasm::SIMDReplace *expr) {
return (std::max)(visit(expr->vec), visit(expr->value));
}
SideEffect visitSIMDShuffle(wasm::SIMDShuffle *expr) {
return (std::max)(visit(expr->left), visit(expr->right));
}
SideEffect visitSIMDBitselect(wasm::SIMDBitselect *expr) {
return (std::max)({
visit(expr->cond),
visit(expr->left),
visit(expr->right),
});
}
SideEffect visitSIMDShift(wasm::SIMDShift *expr) {
return (std::max)(visit(expr->vec), visit(expr->shift));
}
SideEffect visitMemoryInit(wasm::MemoryInit *expr) {
visit(expr->dest);
visit(expr->offset);
visit(expr->size);
return SideEffect::write;
}
SideEffect visitDataDrop([[maybe_unused]] wasm::DataDrop *expr) {
return SideEffect::write;
}
SideEffect visitMemoryCopy(wasm::MemoryCopy *expr) {
visit(expr->dest);
visit(expr->source);
visit(expr->size);
return SideEffect::write;
}
SideEffect visitMemoryFill(wasm::MemoryFill *expr) {
visit(expr->dest);
visit(expr->value);
visit(expr->size);
return SideEffect::write;
}
SideEffect visit(wasm::Expression *expr) {
if (expr == nullptr) {
return SideEffect::none;
}
return OverriddenVisitor::visit(expr);
} }
void visitFunction(wasm::Function *function) { void visitFunction(wasm::Function *function) {
...@@ -842,10 +586,19 @@ namespace kyut::watermarker { ...@@ -842,10 +586,19 @@ namespace kyut::watermarker {
// Embed watermarks // Embed watermarks
const auto posStart = stream.tell(); const auto posStart = stream.tell();
EmbeddingVisitor embedder{stream};
OperandSwappingVisitor visitor{[&stream](wasm::Binary &expr) {
// Embed watermarks by swapping operands
const bool bit = stream.readBit();
const auto &lo = (std::min)(*expr.left, *expr.right);
if (bit == (expr.left == &lo)) {
swapOperands(expr);
}
}};
for (const auto f : functions) { for (const auto f : functions) {
embedder.visitFunction(f); visitor.visitFunction(f);
} }
return stream.tell() - posStart; return stream.tell() - posStart;
...@@ -866,10 +619,16 @@ namespace kyut::watermarker { ...@@ -866,10 +619,16 @@ namespace kyut::watermarker {
// Extract watermarks // Extract watermarks
const auto posStart = stream.tell(); const auto posStart = stream.tell();
ExtractingVisitor extractor{stream};
OperandSwappingVisitor visitor{[&stream](wasm::Binary &expr) {
// Extract watermarks from the order of operands
const auto &lo = (std::min)(*expr.left, *expr.right);
stream.writeBit(expr.left != &lo);
}};
for (const auto f : functions) { for (const auto f : functions) {
extractor.visitFunction(f); visitor.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