Add a utility class for extracting details of the DOM

headless/public/util/dom_tree_extractor.cc

+// 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