Add method to get ScriptBase data without creating an AnyScript

citp · May 5, 2018 · 44ffee4 · 44ffee4
1 parent c2778be
commit 44ffee4
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 2 deletions.
diff --git a/include/blocksci/address/address.hpp b/include/blocksci/address/address.hpp
@@ -63,6 +63,8 @@ namespace blocksci {
 
         AnyScript getScript() const;
 
+        ScriptBase getBaseScript() const;
+
         EquivAddress getEquivAddresses(bool nestedEquivalent) const;
 
         ranges::any_view<OutputPointer> getOutputPointers() const;

diff --git a/include/blocksci/scripts/script.hpp b/include/blocksci/scripts/script.hpp
@@ -27,6 +27,8 @@ namespace blocksci {
         ScriptBase() = default;
         ScriptBase(uint32_t scriptNum_, AddressType::Enum type_, DataAccess &access_, const ScriptDataBase *rawData_) : Address(scriptNum_, type_, access_), rawData(rawData_) {}
 
+        explicit ScriptBase(const Address &address);
+
         void visitPointers(const std::function<void(const Address &)> &) const {}
 
         uint32_t getFirstTxIndex() const {

diff --git a/include/blocksci/scripts/script_access.hpp b/include/blocksci/scripts/script_access.hpp
@@ -44,6 +44,11 @@ namespace blocksci {
     template<DedupAddressType::Enum type>
     using ScriptFilePtr = std::unique_ptr<ScriptFile<type>>;
 
+    template<DedupAddressType::Enum type>
+    struct ScriptDataBaseFunctor {
+        static const ScriptDataBase * f(uint32_t scriptNum, const ScriptAccess &access);
+    };
+
     class BLOCKSCI_EXPORT ScriptAccess {
     private:
         using ScriptFilesTuple = to_dedup_address_tuple_t<ScriptFilePtr>;
@@ -67,6 +72,15 @@ namespace blocksci {
             return getFile<type>()[addressNum - 1];
         }
 
+        const ScriptDataBase *getScriptHeader(uint32_t addressNum, DedupAddressType::Enum type) const {
+            static auto &scriptDataBaseTable = *[]() {
+                auto table = make_dynamic_table<DedupAddressType, ScriptDataBaseFunctor>();
+                return new decltype(table){table};
+            }();
+            auto index = static_cast<size_t>(type);
+            return scriptDataBaseTable.at(index)(addressNum, *this);
+        }
+
         std::array<uint32_t, DedupAddressType::size> scriptCounts() const;
 
         uint32_t scriptCount(DedupAddressType::Enum type) const;
@@ -78,6 +92,12 @@ namespace blocksci {
         }
     };
 
+    template<DedupAddressType::Enum type>
+    const ScriptDataBase * ScriptDataBaseFunctor<type>::f(uint32_t scriptNum, const ScriptAccess &access) {
+        auto &file = access.getFile<type>();
+        return file.getDataAtIndex(scriptNum - 1);
+    }
+
 } // namespace blocksci
 
 #endif /* script_access_hpp */
diff --git a/src/address/address.cpp b/src/address/address.cpp
@@ -49,6 +49,10 @@ namespace blocksci {
         return AnyScript{*this};
     }
 
+    ScriptBase Address::getBaseScript() const {
+        return ScriptBase(*this);
+    }
+
     ranges::optional<Address> getAddressFromString(const std::string &addressString, DataAccess &access) {
         if (addressString.compare(0, access.config.segwitPrefix.size(), access.config.segwitPrefix) == 0) {
             std::pair<int, std::vector<uint8_t> > decoded = segwit_addr::decode(access.config.segwitPrefix, addressString);

diff --git a/src/heuristics/change_address.cpp b/src/heuristics/change_address.cpp
@@ -209,7 +209,7 @@ namespace blocksci { namespace heuristics {
         std::unordered_set<Output> candidates;
 
         for (auto output : tx.outputs()) {
-            if (output.getAddress().isSpendable() && output.getAddress().getScript().firstTxIndex() == tx.txNum) {
+            if (output.getAddress().isSpendable() && output.getAddress().getBaseScript().getFirstTxIndex() == tx.txNum) {
                 candidates.insert(output);
             }
         }
@@ -232,7 +232,7 @@ namespace blocksci { namespace heuristics {
         for (auto output : tx.outputs()) {
             if (output.getAddress().isSpendable()) {
                 spendableCount++;
-                if (output.getValue() < smallestInput && output.getAddress().getScript().firstTxIndex() == tx.txNum) {
+                if (output.getValue() < smallestInput && output.getAddress().getBaseScript().getFirstTxIndex() == tx.txNum) {
                     if (change) {
                         return ranges::nullopt;
                     }

diff --git a/src/scripts/script.cpp b/src/scripts/script.cpp
@@ -6,9 +6,14 @@
 //
 
 #include <blocksci/scripts/script.hpp>
+#include <blocksci/scripts/script_access.hpp>
 #include <blocksci/chain/transaction.hpp>
+#include <blocksci/core/address_info.hpp>
 
 namespace blocksci {
+
+    ScriptBase::ScriptBase(const Address &address) : ScriptBase(address.scriptNum, address.type, address.getAccess(), address.getAccess().scripts->getScriptHeader(address.scriptNum, dedupType(address.type))) {}
+
     Transaction ScriptBase::getFirstTransaction() const {
         return Transaction(getFirstTxIndex(), getAccess());
     }