Add block height start and stop to create_clustering, Fixes #118

citp · Oct 22, 2018 · 5b1584d · 5b1584d
1 parent 2910bcb
commit 5b1584d
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 9 deletions.
diff --git a/blockscipy/src/cluster/cluster/cluster_py.cpp b/blockscipy/src/cluster/cluster/cluster_py.cpp
@@ -53,13 +53,17 @@ void init_cluster_manager(pybind11::module &s) {
     .def(py::init([](std::string arg, blocksci::Blockchain &chain) {
        return ClusterManager(arg, chain.getAccess());
     }))
-    .def_static("create_clustering", [](const std::string &location, Blockchain &chain, const heuristics::ChangeHeuristic &heuristic, bool shouldOverwrite) {
+    .def_static("create_clustering", [](const std::string &location, Blockchain &chain, BlockHeight start, BlockHeight stop, const heuristics::ChangeHeuristic &heuristic, bool shouldOverwrite) {
         py::scoped_ostream_redirect stream(
                                            std::cout,
                                            py::module::import("sys").attr("stdout")
                                            );
-        return ClusterManager::createClustering(chain, heuristic, location, shouldOverwrite);
-    }, py::arg("location"), py::arg("chain"), py::arg("heuristic") = heuristics::ChangeHeuristic{heuristics::LegacyChange{}}, py::arg("should_overwrite") = false)
+        if (stop == -1) {
+            stop = chain.size();
+        }
+        auto range = chain[{start, stop}];
+        return ClusterManager::createClustering(range, heuristic, location, shouldOverwrite);
+    }, py::arg("location"), py::arg("chain"), py::arg("start") = 0, py::arg("stop") = -1, py::arg("heuristic") = heuristics::ChangeHeuristic{heuristics::LegacyChange{}}, py::arg("should_overwrite") = false)
     .def("cluster_with_address", [](const ClusterManager &cm, const Address &address) -> Cluster {
        return cm.getCluster(address);
     }, py::arg("address"), "Return the cluster containing the given address")

diff --git a/include/blocksci/chain/block_range.hpp b/include/blocksci/chain/block_range.hpp
@@ -197,6 +197,9 @@ namespace blocksci {
         std::vector<BlockRange> segment(unsigned int segmentCount) const;
 
         Slice sl;
+
+        DataAccess &getAccess() { return *access; }
+
     private:
         DataAccess *access;
 

diff --git a/include/blocksci/chain/blockchain.hpp b/include/blocksci/chain/blockchain.hpp
@@ -36,8 +36,6 @@ namespace blocksci {
 
         void reload();
 
-        DataAccess &getAccess() { return *access; }
-
         uint32_t addressCount(AddressType::Enum type) const;
     };
 

diff --git a/include/blocksci/cluster/cluster_manager.hpp b/include/blocksci/cluster/cluster_manager.hpp
@@ -33,7 +33,7 @@ namespace blocksci {
         ClusterManager &operator=(ClusterManager && other);
         ~ClusterManager();
 
-        static ClusterManager createClustering(Blockchain &chain, const heuristics::ChangeHeuristic &heuristic, const std::string &outputPath, bool overwrite = false);
+        static ClusterManager createClustering(BlockRange &chain, const heuristics::ChangeHeuristic &heuristic, const std::string &outputPath, bool overwrite = false);
 
         Cluster getCluster(const Address &address) const;
 

diff --git a/src/cluster/cluster_manager.cpp b/src/cluster/cluster_manager.cpp
@@ -149,13 +149,13 @@ namespace blocksci {
         }
     };
 
-    std::vector<uint32_t> createClusters(Blockchain &chain, std::unordered_map<DedupAddressType::Enum, uint32_t> addressStarts, uint32_t totalScriptCount, const heuristics::ChangeHeuristic &changeHeuristic) {
+    std::vector<uint32_t> createClusters(BlockRange &chain, std::unordered_map<DedupAddressType::Enum, uint32_t> addressStarts, uint32_t totalScriptCount, const heuristics::ChangeHeuristic &changeHeuristic) {
 
         AddressDisjointSets ds(totalScriptCount, std::move(addressStarts));
 
         auto &access = chain.getAccess();
 
-        auto scriptHashCount = chain.addressCount(AddressType::SCRIPTHASH);
+        auto scriptHashCount = access.getScripts().scriptCount(DedupAddressType::SCRIPTHASH);
 
 
         segmentWork(1, scriptHashCount + 1, 8, [&ds, &access](uint32_t index) {
@@ -246,7 +246,7 @@ namespace blocksci {
     }
 
 
-    ClusterManager ClusterManager::createClustering(Blockchain &chain, const heuristics::ChangeHeuristic &changeHeuristic, const std::string &outputPath, bool overwrite) {
+    ClusterManager ClusterManager::createClustering(BlockRange &chain, const heuristics::ChangeHeuristic &changeHeuristic, const std::string &outputPath, bool overwrite) {
 
         auto outputLocation = filesystem::path{outputPath};