Index: headless/public/util/dom_tree_extractor.cc |
diff --git a/headless/public/util/dom_tree_extractor.cc b/headless/public/util/dom_tree_extractor.cc |
new file mode 100644 |
index 0000000000000000000000000000000000000000..70ed8afe046a2bf87153f00bfb89bd11207ecd80 |
--- /dev/null |
+++ b/headless/public/util/dom_tree_extractor.cc |
@@ -0,0 +1,115 @@ |
+// Copyright 2016 The Chromium Authors. All rights reserved. |
+// Use of this source code is governed by a BSD-style license that can be |
+// found in the LICENSE file. |
+ |
+#include "headless/public/util/dom_tree_extractor.h" |
+ |
+#include "base/bind.h" |
+#include "base/json/json_writer.h" |
+#include "headless/public/headless_devtools_client.h" |
+ |
+namespace headless { |
+ |
+DomTreeExtractor::DomTreeExtractor(HeadlessDevToolsClient* devtools_client) |
+ : work_in_progress_(false), |
+ devtools_client_(devtools_client), |
+ weak_factory_(this) {} |
+ |
+DomTreeExtractor::~DomTreeExtractor() {} |
+ |
+void DomTreeExtractor::ExtractDomTree( |
+ const std::vector<std::string>& css_style_whitelist, |
+ DomResultCB callback) { |
+ DCHECK(!work_in_progress_); |
+ work_in_progress_ = true; |
+ |
+ callback_ = std::move(callback); |
+ |
+ devtools_client_->GetDOM()->GetDocument( |
+ dom::GetDocumentParams::Builder() |
+ .SetDepth(-1) |
+ .SetTraverseFrames(true) |
+ .Build(), |
+ base::Bind(&DomTreeExtractor::OnDocumentFetched, |
+ weak_factory_.GetWeakPtr())); |
+ |
+ devtools_client_->GetCSS()->GetExperimental()->GetLayoutTreeAndStyles( |
+ css::GetLayoutTreeAndStylesParams::Builder() |
+ .SetComputedStyleWhitelist(css_style_whitelist) |
+ .Build(), |
+ base::Bind(&DomTreeExtractor::OnLayoutTreeAndStylesFetched, |
+ weak_factory_.GetWeakPtr())); |
+} |
+ |
+void DomTreeExtractor::OnDocumentFetched( |
+ std::unique_ptr<dom::GetDocumentResult> result) { |
+ dom_tree_.document_result_ = std::move(result); |
+ MaybeExtractDomTree(); |
+} |
+ |
+void DomTreeExtractor::OnLayoutTreeAndStylesFetched( |
+ std::unique_ptr<css::GetLayoutTreeAndStylesResult> result) { |
+ dom_tree_.layout_tree_and_styles_result_ = std::move(result); |
+ MaybeExtractDomTree(); |
+} |
+ |
+void DomTreeExtractor::MaybeExtractDomTree() { |
+ if (dom_tree_.document_result_ && dom_tree_.layout_tree_and_styles_result_) { |
+ EnumerateNodes(dom_tree_.document_result_->GetRoot()); |
+ ExtractLayoutTreeNodes(); |
+ ExtractComputedStyles(); |
+ |
+ work_in_progress_ = false; |
+ |
+ callback_.Run(std::move(dom_tree_)); |
+ } |
+} |
+ |
+void DomTreeExtractor::EnumerateNodes(const dom::Node* node) { |
+ // Allocate an index and record the node pointer. |
+ size_t index = dom_tree_.node_id_to_index_.size(); |
+ dom_tree_.node_id_to_index_[node->GetNodeId()] = index; |
+ dom_tree_.dom_nodes_.push_back(node); |
+ |
+ if (node->HasContentDocument()) |
+ EnumerateNodes(node->GetContentDocument()); |
+ |
+ if (node->HasChildren()) { |
+ for (const std::unique_ptr<dom::Node>& child : *node->GetChildren()) { |
+ EnumerateNodes(child.get()); |
+ } |
+ } |
+} |
+ |
+void DomTreeExtractor::ExtractLayoutTreeNodes() { |
+ dom_tree_.layout_tree_nodes_.reserve( |
+ dom_tree_.layout_tree_and_styles_result_->GetLayoutTreeNodes()->size()); |
+ |
+ // Only extract layout tree nodes that map to a DOM node. |
+ for (const std::unique_ptr<css::LayoutTreeNode>& layout_node : |
+ *dom_tree_.layout_tree_and_styles_result_->GetLayoutTreeNodes()) { |
+ std::unordered_map<NodeId, size_t>::const_iterator it = |
+ dom_tree_.node_id_to_index_.find(layout_node->GetBackendNodeId()); |
+ if (it == dom_tree_.node_id_to_index_.end()) |
+ continue; |
+ |
+ dom_tree_.layout_tree_nodes_.push_back(layout_node.get()); |
+ } |
+} |
+ |
+void DomTreeExtractor::ExtractComputedStyles() { |
+ dom_tree_.computed_styles_.reserve( |
+ dom_tree_.layout_tree_and_styles_result_->GetComputedStyles()->size()); |
+ |
+ for (const std::unique_ptr<css::ComputedStyle>& computed_style : |
+ *dom_tree_.layout_tree_and_styles_result_->GetComputedStyles()) { |
+ dom_tree_.computed_styles_.push_back(computed_style.get()); |
+ } |
+} |
+ |
+DomTreeExtractor::DomTree::DomTree() {} |
+DomTreeExtractor::DomTree::~DomTree() {} |
+ |
+DomTreeExtractor::DomTree::DomTree(DomTree&& other) = default; |
+ |
+} // namespace headless |