AAPT2: Change XmlDom to exclude Namespace as a node

In preparation for exporting an XML proto format for UAM to consume,
this change brings the XML DOM API more in line with other APIs that
do not make the Namespace a separate node.

Treating Namespace declarations as just properties of an Element
node makes the implementation of algorithms much simpler, as
the constraints that Namespace nodes have only one child
are now built in and traversing to find Element nodes
is much simpler.

Also made a bunch of quality of life improvements, like formatting and
comment style.

Test: make aapt2_tests
Change-Id: Ib97ff1c4252b7907e2cc1f13a448dc4ca3b809a4
diff --git a/tools/aapt2/xml/XmlActionExecutor.cpp b/tools/aapt2/xml/XmlActionExecutor.cpp
index 352a933..cc664a5 100644
--- a/tools/aapt2/xml/XmlActionExecutor.cpp
+++ b/tools/aapt2/xml/XmlActionExecutor.cpp
@@ -78,7 +78,7 @@
                                 XmlResource* doc) const {
   SourcePathDiagnostics source_diag(doc->file.source, diag);
 
-  Element* el = FindRootElement(doc);
+  Element* el = doc->root.get();
   if (!el) {
     if (policy == XmlActionExecutorPolicy::kWhitelist) {
       source_diag.Error(DiagMessage() << "no root XML tag found");
diff --git a/tools/aapt2/xml/XmlDom.cpp b/tools/aapt2/xml/XmlDom.cpp
index d6df715..cbb652e 100644
--- a/tools/aapt2/xml/XmlDom.cpp
+++ b/tools/aapt2/xml/XmlDom.cpp
@@ -39,17 +39,15 @@
 constexpr char kXmlNamespaceSep = 1;
 
 struct Stack {
-  std::unique_ptr<xml::Node> root;
-  std::stack<xml::Node*> node_stack;
+  std::unique_ptr<xml::Element> root;
+  std::stack<xml::Element*> node_stack;
+  std::unique_ptr<xml::Element> pending_element;
   std::string pending_comment;
   std::unique_ptr<xml::Text> last_text_node;
 };
 
-/**
- * Extracts the namespace and name of an expanded element or attribute name.
- */
-static void SplitName(const char* name, std::string* out_ns,
-                      std::string* out_name) {
+// Extracts the namespace and name of an expanded element or attribute name.
+static void SplitName(const char* name, std::string* out_ns, std::string* out_name) {
   const char* p = name;
   while (*p != 0 && *p != kXmlNamespaceSep) {
     p++;
@@ -67,6 +65,7 @@
 static void FinishPendingText(Stack* stack) {
   if (stack->last_text_node != nullptr) {
     if (!stack->last_text_node->text.empty()) {
+      CHECK(!stack->node_stack.empty());
       stack->node_stack.top()->AppendChild(std::move(stack->last_text_node));
     } else {
       // Drop an empty text node.
@@ -75,48 +74,27 @@
   }
 }
 
-static void AddToStack(Stack* stack, XML_Parser parser,
-                       std::unique_ptr<Node> node) {
-  node->line_number = XML_GetCurrentLineNumber(parser);
-  node->column_number = XML_GetCurrentColumnNumber(parser);
-
-  Node* this_node = node.get();
-  if (!stack->node_stack.empty()) {
-    stack->node_stack.top()->AppendChild(std::move(node));
-  } else {
-    stack->root = std::move(node);
-  }
-
-  if (!NodeCast<Text>(this_node)) {
-    stack->node_stack.push(this_node);
-  }
-}
-
-static void XMLCALL StartNamespaceHandler(void* user_data, const char* prefix,
-                                          const char* uri) {
+static void XMLCALL StartNamespaceHandler(void* user_data, const char* prefix, const char* uri) {
   XML_Parser parser = reinterpret_cast<XML_Parser>(user_data);
   Stack* stack = reinterpret_cast<Stack*>(XML_GetUserData(parser));
   FinishPendingText(stack);
 
-  std::unique_ptr<Namespace> ns = util::make_unique<Namespace>();
-  if (prefix) {
-    ns->namespace_prefix = prefix;
-  }
+  NamespaceDecl decl;
+  decl.line_number = XML_GetCurrentLineNumber(parser);
+  decl.column_number = XML_GetCurrentColumnNumber(parser);
+  decl.prefix = prefix ? prefix : "";
+  decl.uri = uri ? uri : "";
 
-  if (uri) {
-    ns->namespace_uri = uri;
+  if (stack->pending_element == nullptr) {
+    stack->pending_element = util::make_unique<Element>();
   }
-
-  AddToStack(stack, parser, std::move(ns));
+  stack->pending_element->namespace_decls.push_back(std::move(decl));
 }
 
-static void XMLCALL EndNamespaceHandler(void* user_data, const char* prefix) {
+static void XMLCALL EndNamespaceHandler(void* user_data, const char* /*prefix*/) {
   XML_Parser parser = reinterpret_cast<XML_Parser>(user_data);
   Stack* stack = reinterpret_cast<Stack*>(XML_GetUserData(parser));
   FinishPendingText(stack);
-
-  CHECK(!stack->node_stack.empty());
-  stack->node_stack.pop();
 }
 
 static bool less_attribute(const Attribute& lhs, const Attribute& rhs) {
@@ -124,28 +102,42 @@
          std::tie(rhs.namespace_uri, rhs.name, rhs.value);
 }
 
-static void XMLCALL StartElementHandler(void* user_data, const char* name,
-                                        const char** attrs) {
+static void XMLCALL StartElementHandler(void* user_data, const char* name, const char** attrs) {
   XML_Parser parser = reinterpret_cast<XML_Parser>(user_data);
   Stack* stack = reinterpret_cast<Stack*>(XML_GetUserData(parser));
   FinishPendingText(stack);
 
-  std::unique_ptr<Element> el = util::make_unique<Element>();
+  std::unique_ptr<Element> el;
+  if (stack->pending_element != nullptr) {
+    el = std::move(stack->pending_element);
+  } else {
+    el = util::make_unique<Element>();
+  }
+
+  el->line_number = XML_GetCurrentLineNumber(parser);
+  el->column_number = XML_GetCurrentColumnNumber(parser);
+  el->comment = std::move(stack->pending_comment);
+
   SplitName(name, &el->namespace_uri, &el->name);
 
   while (*attrs) {
     Attribute attribute;
     SplitName(*attrs++, &attribute.namespace_uri, &attribute.name);
     attribute.value = *attrs++;
-
-    // Insert in sorted order.
-    auto iter = std::lower_bound(el->attributes.begin(), el->attributes.end(), attribute,
-                                 less_attribute);
-    el->attributes.insert(iter, std::move(attribute));
+    el->attributes.push_back(std::move(attribute));
   }
 
-  el->comment = std::move(stack->pending_comment);
-  AddToStack(stack, parser, std::move(el));
+  // Sort the attributes.
+  std::sort(el->attributes.begin(), el->attributes.end(), less_attribute);
+
+  // Add to the stack.
+  Element* this_el = el.get();
+  if (!stack->node_stack.empty()) {
+    stack->node_stack.top()->AppendChild(std::move(el));
+  } else {
+    stack->root = std::move(el);
+  }
+  stack->node_stack.push(this_el);
 }
 
 static void XMLCALL EndElementHandler(void* user_data, const char* name) {
@@ -263,13 +255,13 @@
 std::unique_ptr<XmlResource> Inflate(const void* data, size_t data_len, IDiagnostics* diag,
                                      const Source& source) {
   // We import the android namespace because on Windows NO_ERROR is a macro, not
-  // an enum, which
-  // causes errors when qualifying it with android::
+  // an enum, which causes errors when qualifying it with android::
   using namespace android;
 
   StringPool string_pool;
-  std::unique_ptr<Node> root;
-  std::stack<Node*> node_stack;
+  std::unique_ptr<Element> root;
+  std::stack<Element*> node_stack;
+  std::unique_ptr<Element> pending_element;
 
   ResXMLTree tree;
   if (tree.setTo(data, data_len) != NO_ERROR) {
@@ -277,57 +269,76 @@
   }
 
   ResXMLParser::event_code_t code;
-  while ((code = tree.next()) != ResXMLParser::BAD_DOCUMENT &&
-         code != ResXMLParser::END_DOCUMENT) {
+  while ((code = tree.next()) != ResXMLParser::BAD_DOCUMENT && code != ResXMLParser::END_DOCUMENT) {
     std::unique_ptr<Node> new_node;
     switch (code) {
       case ResXMLParser::START_NAMESPACE: {
-        std::unique_ptr<Namespace> node = util::make_unique<Namespace>();
+        NamespaceDecl decl;
         size_t len;
         const char16_t* str16 = tree.getNamespacePrefix(&len);
         if (str16) {
-          node->namespace_prefix = util::Utf16ToUtf8(StringPiece16(str16, len));
+          decl.prefix = util::Utf16ToUtf8(StringPiece16(str16, len));
         }
 
         str16 = tree.getNamespaceUri(&len);
         if (str16) {
-          node->namespace_uri = util::Utf16ToUtf8(StringPiece16(str16, len));
+          decl.uri = util::Utf16ToUtf8(StringPiece16(str16, len));
         }
-        new_node = std::move(node);
+
+        if (pending_element == nullptr) {
+          pending_element = util::make_unique<Element>();
+        }
         break;
       }
 
       case ResXMLParser::START_TAG: {
-        std::unique_ptr<Element> node = util::make_unique<Element>();
+        std::unique_ptr<Element> el;
+        if (pending_element != nullptr) {
+          el = std::move(pending_element);
+        } else {
+          el = util::make_unique<Element>();
+          ;
+        }
+
         size_t len;
         const char16_t* str16 = tree.getElementNamespace(&len);
         if (str16) {
-          node->namespace_uri = util::Utf16ToUtf8(StringPiece16(str16, len));
+          el->namespace_uri = util::Utf16ToUtf8(StringPiece16(str16, len));
         }
 
         str16 = tree.getElementName(&len);
         if (str16) {
-          node->name = util::Utf16ToUtf8(StringPiece16(str16, len));
+          el->name = util::Utf16ToUtf8(StringPiece16(str16, len));
         }
 
-        CopyAttributes(node.get(), &tree, &string_pool);
+        Element* this_el = el.get();
+        CopyAttributes(el.get(), &tree, &string_pool);
 
-        new_node = std::move(node);
+        if (!node_stack.empty()) {
+          node_stack.top()->AppendChild(std::move(el));
+        } else {
+          root = std::move(el);
+        }
+        node_stack.push(this_el);
         break;
       }
 
       case ResXMLParser::TEXT: {
-        std::unique_ptr<Text> node = util::make_unique<Text>();
+        std::unique_ptr<Text> text = util::make_unique<Text>();
+        text->line_number = tree.getLineNumber();
         size_t len;
         const char16_t* str16 = tree.getText(&len);
         if (str16) {
-          node->text = util::Utf16ToUtf8(StringPiece16(str16, len));
+          text->text = util::Utf16ToUtf8(StringPiece16(str16, len));
         }
-        new_node = std::move(node);
+        CHECK(!node_stack.empty());
+        node_stack.top()->AppendChild(std::move(text));
         break;
       }
 
       case ResXMLParser::END_NAMESPACE:
+        break;
+
       case ResXMLParser::END_TAG:
         CHECK(!node_stack.empty());
         node_stack.pop();
@@ -337,74 +348,32 @@
         LOG(FATAL) << "unhandled XML chunk type";
         break;
     }
-
-    if (new_node) {
-      new_node->line_number = tree.getLineNumber();
-
-      Node* this_node = new_node.get();
-      if (!root) {
-        CHECK(node_stack.empty()) << "node stack should be empty";
-        root = std::move(new_node);
-      } else {
-        CHECK(!node_stack.empty()) << "node stack should not be empty";
-        node_stack.top()->AppendChild(std::move(new_node));
-      }
-
-      if (!NodeCast<Text>(this_node)) {
-        node_stack.push(this_node);
-      }
-    }
   }
   return util::make_unique<XmlResource>(ResourceFile{}, std::move(string_pool), std::move(root));
 }
 
-std::unique_ptr<Node> Namespace::Clone(const ElementCloneFunc& el_cloner) {
-  auto ns = util::make_unique<Namespace>();
-  ns->comment = comment;
-  ns->line_number = line_number;
-  ns->column_number = column_number;
-  ns->namespace_prefix = namespace_prefix;
-  ns->namespace_uri = namespace_uri;
-  ns->children.reserve(children.size());
-  for (const std::unique_ptr<xml::Node>& child : children) {
-    ns->AppendChild(child->Clone(el_cloner));
-  }
-  return std::move(ns);
-}
-
-Element* FindRootElement(XmlResource* doc) {
-  return FindRootElement(doc->root.get());
-}
-
 Element* FindRootElement(Node* node) {
-  if (!node) {
+  if (node == nullptr) {
     return nullptr;
   }
 
-  Element* el = nullptr;
-  while ((el = NodeCast<Element>(node)) == nullptr) {
-    if (node->children.empty()) {
-      return nullptr;
-    }
-    // We are looking for the first element, and namespaces can only have one
-    // child.
-    node = node->children.front().get();
+  while (node->parent != nullptr) {
+    node = node->parent;
   }
-  return el;
+  return NodeCast<Element>(node);
 }
 
-void Node::AppendChild(std::unique_ptr<Node> child) {
+void Element::AppendChild(std::unique_ptr<Node> child) {
   child->parent = this;
   children.push_back(std::move(child));
 }
 
-void Node::InsertChild(size_t index, std::unique_ptr<Node> child) {
+void Element::InsertChild(size_t index, std::unique_ptr<Node> child) {
   child->parent = this;
   children.insert(children.begin() + index, std::move(child));
 }
 
-Attribute* Element::FindAttribute(const StringPiece& ns,
-                                  const StringPiece& name) {
+Attribute* Element::FindAttribute(const StringPiece& ns, const StringPiece& name) {
   for (auto& attr : attributes) {
     if (ns == attr.namespace_uri && name == attr.name) {
       return &attr;
@@ -426,21 +395,11 @@
   return FindChildWithAttribute(ns, name, {}, {}, {});
 }
 
-Element* Element::FindChildWithAttribute(const StringPiece& ns,
-                                         const StringPiece& name,
-                                         const StringPiece& attr_ns,
-                                         const StringPiece& attr_name,
+Element* Element::FindChildWithAttribute(const StringPiece& ns, const StringPiece& name,
+                                         const StringPiece& attr_ns, const StringPiece& attr_name,
                                          const StringPiece& attr_value) {
-  for (auto& child_node : children) {
-    Node* child = child_node.get();
-    while (NodeCast<Namespace>(child)) {
-      if (child->children.empty()) {
-        break;
-      }
-      child = child->children[0].get();
-    }
-
-    if (Element* el = NodeCast<Element>(child)) {
+  for (auto& child : children) {
+    if (Element* el = NodeCast<Element>(child.get())) {
       if (ns == el->namespace_uri && name == el->name) {
         if (attr_ns.empty() && attr_name.empty()) {
           return el;
@@ -459,23 +418,16 @@
 std::vector<Element*> Element::GetChildElements() {
   std::vector<Element*> elements;
   for (auto& child_node : children) {
-    Node* child = child_node.get();
-    while (NodeCast<Namespace>(child)) {
-      if (child->children.empty()) {
-        break;
-      }
-      child = child->children[0].get();
-    }
-
-    if (Element* el = NodeCast<Element>(child)) {
-      elements.push_back(el);
+    if (Element* child = NodeCast<Element>(child_node.get())) {
+      elements.push_back(child);
     }
   }
   return elements;
 }
 
-std::unique_ptr<Node> Element::Clone(const ElementCloneFunc& el_cloner) {
+std::unique_ptr<Node> Element::Clone(const ElementCloneFunc& el_cloner) const {
   auto el = util::make_unique<Element>();
+  el->namespace_decls = namespace_decls;
   el->comment = comment;
   el->line_number = line_number;
   el->column_number = column_number;
@@ -490,7 +442,17 @@
   return std::move(el);
 }
 
-std::unique_ptr<Node> Text::Clone(const ElementCloneFunc&) {
+std::unique_ptr<Element> Element::CloneElement(const ElementCloneFunc& el_cloner) const {
+  return std::unique_ptr<Element>(static_cast<Element*>(Clone(el_cloner).release()));
+}
+
+void Element::Accept(Visitor* visitor) {
+  visitor->BeforeVisitElement(this);
+  visitor->Visit(this);
+  visitor->AfterVisitElement(this);
+}
+
+std::unique_ptr<Node> Text::Clone(const ElementCloneFunc&) const {
   auto t = util::make_unique<Text>();
   t->comment = comment;
   t->line_number = line_number;
@@ -499,21 +461,22 @@
   return std::move(t);
 }
 
-void PackageAwareVisitor::Visit(Namespace* ns) {
-  bool added = false;
-  if (Maybe<ExtractedPackage> maybe_package =
-          ExtractPackageFromNamespace(ns->namespace_uri)) {
-    ExtractedPackage& package = maybe_package.value();
-    package_decls_.push_back(
-        PackageDecl{ns->namespace_prefix, std::move(package)});
-    added = true;
-  }
+void Text::Accept(Visitor* visitor) {
+  visitor->Visit(this);
+}
 
-  Visitor::Visit(ns);
-
-  if (added) {
-    package_decls_.pop_back();
+void PackageAwareVisitor::BeforeVisitElement(Element* el) {
+  std::vector<PackageDecl> decls;
+  for (const NamespaceDecl& decl : el->namespace_decls) {
+    if (Maybe<ExtractedPackage> maybe_package = ExtractPackageFromNamespace(decl.uri)) {
+      decls.push_back(PackageDecl{decl.prefix, std::move(maybe_package.value())});
+    }
   }
+  package_decls_.push_back(std::move(decls));
+}
+
+void PackageAwareVisitor::AfterVisitElement(Element* el) {
+  package_decls_.pop_back();
 }
 
 Maybe<ExtractedPackage> PackageAwareVisitor::TransformPackageAlias(
@@ -524,11 +487,16 @@
 
   const auto rend = package_decls_.rend();
   for (auto iter = package_decls_.rbegin(); iter != rend; ++iter) {
-    if (alias == iter->prefix) {
-      if (iter->package.package.empty()) {
-        return ExtractedPackage{local_package.to_string(), iter->package.private_namespace};
+    const std::vector<PackageDecl>& decls = *iter;
+    const auto rend2 = decls.rend();
+    for (auto iter2 = decls.rbegin(); iter2 != rend2; ++iter2) {
+      const PackageDecl& decl = *iter2;
+      if (alias == decl.prefix) {
+        if (decl.package.package.empty()) {
+          return ExtractedPackage{local_package.to_string(), decl.package.private_namespace};
+        }
+        return decl.package;
       }
-      return iter->package;
     }
   }
   return {};
diff --git a/tools/aapt2/xml/XmlDom.h b/tools/aapt2/xml/XmlDom.h
index 54a7033..1542243 100644
--- a/tools/aapt2/xml/XmlDom.h
+++ b/tools/aapt2/xml/XmlDom.h
@@ -33,45 +33,33 @@
 namespace aapt {
 namespace xml {
 
-class RawVisitor;
-
 class Element;
+class Visitor;
 
 // Base class for all XML nodes.
 class Node {
  public:
-  Node* parent = nullptr;
-  size_t line_number = 0;
-  size_t column_number = 0;
-  std::string comment;
-  std::vector<std::unique_ptr<Node>> children;
-
   virtual ~Node() = default;
 
-  void AppendChild(std::unique_ptr<Node> child);
-  void InsertChild(size_t index, std::unique_ptr<Node> child);
-  virtual void Accept(RawVisitor* visitor) = 0;
+  Element* parent = nullptr;
+  size_t line_number = 0u;
+  size_t column_number = 0u;
+  std::string comment;
+
+  virtual void Accept(Visitor* visitor) = 0;
 
   using ElementCloneFunc = std::function<void(const Element&, Element*)>;
 
   // Clones the Node subtree, using the given function to decide how to clone an Element.
-  virtual std::unique_ptr<Node> Clone(const ElementCloneFunc& el_cloner) = 0;
+  virtual std::unique_ptr<Node> Clone(const ElementCloneFunc& el_cloner) const = 0;
 };
 
-// Base class that implements the visitor methods for a subclass of Node.
-template <typename Derived>
-class BaseNode : public Node {
- public:
-  virtual void Accept(RawVisitor* visitor) override;
-};
-
-// A Namespace XML node. Can only have one child.
-class Namespace : public BaseNode<Namespace> {
- public:
-  std::string namespace_prefix;
-  std::string namespace_uri;
-
-  std::unique_ptr<Node> Clone(const ElementCloneFunc& el_cloner) override;
+// A namespace declaration (xmlns:prefix="uri").
+struct NamespaceDecl {
+  std::string prefix;
+  std::string uri;
+  size_t line_number = 0u;
+  size_t column_number = 0u;
 };
 
 struct AaptAttribute {
@@ -94,31 +82,46 @@
 };
 
 // An Element XML node.
-class Element : public BaseNode<Element> {
+class Element : public Node {
  public:
+  // Ordered namespace prefix declarations.
+  std::vector<NamespaceDecl> namespace_decls;
+
   std::string namespace_uri;
   std::string name;
   std::vector<Attribute> attributes;
+  std::vector<std::unique_ptr<Node>> children;
+
+  void AppendChild(std::unique_ptr<Node> child);
+  void InsertChild(size_t index, std::unique_ptr<Node> child);
 
   Attribute* FindAttribute(const android::StringPiece& ns, const android::StringPiece& name);
   const Attribute* FindAttribute(const android::StringPiece& ns,
                                  const android::StringPiece& name) const;
-  xml::Element* FindChild(const android::StringPiece& ns, const android::StringPiece& name);
-  xml::Element* FindChildWithAttribute(const android::StringPiece& ns,
-                                       const android::StringPiece& name,
-                                       const android::StringPiece& attr_ns,
-                                       const android::StringPiece& attr_name,
-                                       const android::StringPiece& attr_value);
-  std::vector<xml::Element*> GetChildElements();
-  std::unique_ptr<Node> Clone(const ElementCloneFunc& el_cloner) override;
+  Element* FindChild(const android::StringPiece& ns, const android::StringPiece& name);
+  Element* FindChildWithAttribute(const android::StringPiece& ns, const android::StringPiece& name,
+                                  const android::StringPiece& attr_ns,
+                                  const android::StringPiece& attr_name,
+                                  const android::StringPiece& attr_value);
+  std::vector<Element*> GetChildElements();
+
+  // Due to overriding of subtypes not working with unique_ptr, define a convenience Clone method
+  // that knows cloning an element returns an element.
+  std::unique_ptr<Element> CloneElement(const ElementCloneFunc& el_cloner) const;
+
+  std::unique_ptr<Node> Clone(const ElementCloneFunc& el_cloner) const override;
+
+  void Accept(Visitor* visitor) override;
 };
 
 // A Text (CDATA) XML node. Can not have any children.
-class Text : public BaseNode<Text> {
+class Text : public Node {
  public:
   std::string text;
 
-  std::unique_ptr<Node> Clone(const ElementCloneFunc& el_cloner) override;
+  std::unique_ptr<Node> Clone(const ElementCloneFunc& el_cloner) const override;
+
+  void Accept(Visitor* visitor) override;
 };
 
 // An XML resource with a source, name, and XML tree.
@@ -131,7 +134,7 @@
   // is destroyed.
   StringPool string_pool;
 
-  std::unique_ptr<xml::Node> root;
+  std::unique_ptr<xml::Element> root;
 };
 
 // Inflates an XML DOM from an InputStream, logging errors to the logger.
@@ -143,42 +146,38 @@
 std::unique_ptr<XmlResource> Inflate(const void* data, size_t data_len, IDiagnostics* diag,
                                      const Source& source);
 
-Element* FindRootElement(XmlResource* doc);
 Element* FindRootElement(Node* node);
 
-// A visitor interface for the different XML Node subtypes. This will not traverse into children.
-// Use Visitor for that.
-class RawVisitor {
- public:
-  virtual ~RawVisitor() = default;
-
-  virtual void Visit(Namespace* node) {}
-  virtual void Visit(Element* node) {}
-  virtual void Visit(Text* text) {}
-};
-
 // Visitor whose default implementation visits the children nodes of any node.
-class Visitor : public RawVisitor {
+class Visitor {
  public:
-  using RawVisitor::Visit;
+  virtual ~Visitor() = default;
 
-  void Visit(Namespace* node) override {
-    VisitChildren(node);
+  virtual void Visit(Element* el) {
+    VisitChildren(el);
   }
 
-  void Visit(Element* node) override {
-    VisitChildren(node);
+  virtual void Visit(Text* text) {
   }
 
-  void Visit(Text* text) override {
-    VisitChildren(text);
-  }
+ protected:
+  Visitor() = default;
 
-  void VisitChildren(Node* node) {
-    for (auto& child : node->children) {
+  void VisitChildren(Element* el) {
+    for (auto& child : el->children) {
       child->Accept(this);
     }
   }
+
+  virtual void BeforeVisitElement(Element* el) {
+  }
+  virtual void AfterVisitElement(Element* el) {
+  }
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(Visitor);
+
+  friend class Element;
 };
 
 // An XML DOM visitor that will record the package name for a namespace prefix.
@@ -186,41 +185,70 @@
  public:
   using Visitor::Visit;
 
-  void Visit(Namespace* ns) override;
   Maybe<ExtractedPackage> TransformPackageAlias(
       const android::StringPiece& alias, const android::StringPiece& local_package) const override;
 
+ protected:
+  PackageAwareVisitor() = default;
+
+  void BeforeVisitElement(Element* el) override;
+  void AfterVisitElement(Element* el) override;
+
  private:
+  DISALLOW_COPY_AND_ASSIGN(PackageAwareVisitor);
+
   struct PackageDecl {
     std::string prefix;
     ExtractedPackage package;
   };
 
-  std::vector<PackageDecl> package_decls_;
+  std::vector<std::vector<PackageDecl>> package_decls_;
 };
 
-// Implementations
+namespace internal {
 
-template <typename Derived>
-void BaseNode<Derived>::Accept(RawVisitor* visitor) {
-  visitor->Visit(static_cast<Derived*>(this));
-}
+// Base class that overrides the default behaviour and does not descend into child nodes.
+class NodeCastBase : public Visitor {
+ public:
+  void Visit(Element* el) override {
+  }
+  void Visit(Text* el) override {
+  }
+
+ protected:
+  NodeCastBase() = default;
+
+  void BeforeVisitElement(Element* el) override {
+  }
+  void AfterVisitElement(Element* el) override {
+  }
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(NodeCastBase);
+};
 
 template <typename T>
-class NodeCastImpl : public RawVisitor {
+class NodeCastImpl : public NodeCastBase {
  public:
-  using RawVisitor::Visit;
+  using NodeCastBase::Visit;
+
+  NodeCastImpl() = default;
 
   T* value = nullptr;
 
   void Visit(T* v) override {
     value = v;
   }
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(NodeCastImpl);
 };
 
+}  // namespace internal
+
 template <typename T>
 T* NodeCast(Node* node) {
-  NodeCastImpl<T> visitor;
+  internal::NodeCastImpl<T> visitor;
   node->Accept(&visitor);
   return visitor.value;
 }
diff --git a/tools/aapt2/xml/XmlDom_test.cpp b/tools/aapt2/xml/XmlDom_test.cpp
index 1340ada..6ed2d61 100644
--- a/tools/aapt2/xml/XmlDom_test.cpp
+++ b/tools/aapt2/xml/XmlDom_test.cpp
@@ -25,8 +25,10 @@
 using ::testing::Eq;
 using ::testing::NotNull;
 using ::testing::SizeIs;
+using ::testing::StrEq;
 
 namespace aapt {
+namespace xml {
 
 TEST(XmlDomTest, Inflate) {
   std::string input = R"(<?xml version="1.0" encoding="utf-8"?>
@@ -40,24 +42,23 @@
 
   StdErrDiagnostics diag;
   StringInputStream in(input);
-  std::unique_ptr<xml::XmlResource> doc = xml::Inflate(&in, &diag, Source("test.xml"));
+  std::unique_ptr<XmlResource> doc = Inflate(&in, &diag, Source("test.xml"));
   ASSERT_THAT(doc, NotNull());
 
-  xml::Namespace* ns = xml::NodeCast<xml::Namespace>(doc->root.get());
-  ASSERT_THAT(ns, NotNull());
-  EXPECT_THAT(ns->namespace_uri, Eq(xml::kSchemaAndroid));
-  EXPECT_THAT(ns->namespace_prefix, Eq("android"));
+  Element* el = doc->root.get();
+  EXPECT_THAT(el->namespace_decls, SizeIs(1u));
+  EXPECT_THAT(el->namespace_decls[0].uri, StrEq(xml::kSchemaAndroid));
+  EXPECT_THAT(el->namespace_decls[0].prefix, StrEq("android"));
 }
 
 // Escaping is handled after parsing of the values for resource-specific values.
 TEST(XmlDomTest, ForwardEscapes) {
-  std::unique_ptr<xml::XmlResource> doc = test::BuildXmlDom(R"(
+  std::unique_ptr<XmlResource> doc = test::BuildXmlDom(R"(
       <element value="\?hello" pattern="\\d{5}">\\d{5}</element>)");
 
-  xml::Element* el = xml::FindRootElement(doc.get());
-  ASSERT_THAT(el, NotNull());
+  Element* el = doc->root.get();
 
-  xml::Attribute* attr = el->FindAttribute({}, "pattern");
+  Attribute* attr = el->FindAttribute({}, "pattern");
   ASSERT_THAT(attr, NotNull());
   EXPECT_THAT(attr->value, Eq("\\\\d{5}"));
 
@@ -67,21 +68,54 @@
 
   ASSERT_THAT(el->children, SizeIs(1u));
 
-  xml::Text* text = xml::NodeCast<xml::Text>(el->children[0].get());
+  Text* text = xml::NodeCast<xml::Text>(el->children[0].get());
   ASSERT_THAT(text, NotNull());
   EXPECT_THAT(text->text, Eq("\\\\d{5}"));
 }
 
 TEST(XmlDomTest, XmlEscapeSequencesAreParsed) {
-  std::unique_ptr<xml::XmlResource> doc = test::BuildXmlDom(R"(<element value="&quot;" />)");
-
-  xml::Element* el = xml::FindRootElement(doc.get());
-  ASSERT_THAT(el, NotNull());
-
-  xml::Attribute* attr = el->FindAttribute({}, "value");
+  std::unique_ptr<XmlResource> doc = test::BuildXmlDom(R"(<element value="&quot;" />)");
+  Attribute* attr = doc->root->FindAttribute({}, "value");
   ASSERT_THAT(attr, NotNull());
-
   EXPECT_THAT(attr->value, Eq("\""));
 }
 
+class TestVisitor : public PackageAwareVisitor {
+ public:
+  using PackageAwareVisitor::Visit;
+
+  void Visit(Element* el) override {
+    if (el->name == "View1") {
+      EXPECT_THAT(TransformPackageAlias("one", "local"),
+                  Eq(make_value(ExtractedPackage{"com.one", false})));
+    } else if (el->name == "View2") {
+      EXPECT_THAT(TransformPackageAlias("one", "local"),
+                  Eq(make_value(ExtractedPackage{"com.one", false})));
+      EXPECT_THAT(TransformPackageAlias("two", "local"),
+                  Eq(make_value(ExtractedPackage{"com.two", false})));
+    } else if (el->name == "View3") {
+      EXPECT_THAT(TransformPackageAlias("one", "local"),
+                  Eq(make_value(ExtractedPackage{"com.one", false})));
+      EXPECT_THAT(TransformPackageAlias("two", "local"),
+                  Eq(make_value(ExtractedPackage{"com.two", false})));
+      EXPECT_THAT(TransformPackageAlias("three", "local"),
+                  Eq(make_value(ExtractedPackage{"com.three", false})));
+    }
+  }
+};
+
+TEST(XmlDomTest, PackageAwareXmlVisitor) {
+  std::unique_ptr<XmlResource> doc = test::BuildXmlDom(R"(
+      <View1 xmlns:one="http://schemas.android.com/apk/res/com.one">
+        <View2 xmlns:two="http://schemas.android.com/apk/res/com.two">
+          <View3 xmlns:three="http://schemas.android.com/apk/res/com.three" />
+        </View2>
+      </View1>)");
+
+  Debug::DumpXml(doc.get());
+  TestVisitor visitor;
+  doc->root->Accept(&visitor);
+}
+
+}  // namespace xml
 }  // namespace aapt
diff --git a/tools/aapt2/xml/XmlUtil.h b/tools/aapt2/xml/XmlUtil.h
index 1650ac2..866b6dc 100644
--- a/tools/aapt2/xml/XmlUtil.h
+++ b/tools/aapt2/xml/XmlUtil.h
@@ -26,79 +26,56 @@
 namespace xml {
 
 constexpr const char* kSchemaAuto = "http://schemas.android.com/apk/res-auto";
-constexpr const char* kSchemaPublicPrefix =
-    "http://schemas.android.com/apk/res/";
-constexpr const char* kSchemaPrivatePrefix =
-    "http://schemas.android.com/apk/prv/res/";
-constexpr const char* kSchemaAndroid =
-    "http://schemas.android.com/apk/res/android";
+constexpr const char* kSchemaPublicPrefix = "http://schemas.android.com/apk/res/";
+constexpr const char* kSchemaPrivatePrefix = "http://schemas.android.com/apk/prv/res/";
+constexpr const char* kSchemaAndroid = "http://schemas.android.com/apk/res/android";
 constexpr const char* kSchemaTools = "http://schemas.android.com/tools";
 constexpr const char* kSchemaAapt = "http://schemas.android.com/aapt";
 
-/**
- * Result of extracting a package name from a namespace URI declaration.
- */
+// Result of extracting a package name from a namespace URI declaration.
 struct ExtractedPackage {
-  /**
-   * The name of the package. This can be the empty string, which means that the
-   * package
-   * should be assumed to be the package being compiled.
-   */
+  // The name of the package. This can be the empty string, which means that the package
+  // should be assumed to be the package being compiled.
   std::string package;
 
-  /**
-   * True if the package's private namespace was declared. This means that
-   * private resources
-   * are made visible.
-   */
+  // True if the package's private namespace was declared. This means that private resources
+  // are made visible.
   bool private_namespace;
+
+  friend inline bool operator==(const ExtractedPackage& a, const ExtractedPackage& b) {
+    return a.package == b.package && a.private_namespace == b.private_namespace;
+  }
 };
 
-/**
- * Returns an ExtractedPackage struct if the namespace URI is of the form:
- * http://schemas.android.com/apk/res/<package> or
- * http://schemas.android.com/apk/prv/res/<package>
- *
- * Special case: if namespaceUri is http://schemas.android.com/apk/res-auto,
- * returns an empty package name.
- */
-Maybe<ExtractedPackage> ExtractPackageFromNamespace(
-    const std::string& namespace_uri);
+// Returns an ExtractedPackage struct if the namespace URI is of the form:
+//   http://schemas.android.com/apk/res/<package> or
+//   http://schemas.android.com/apk/prv/res/<package>
+//
+// Special case: if namespaceUri is http://schemas.android.com/apk/res-auto,
+// returns an empty package name.
+Maybe<ExtractedPackage> ExtractPackageFromNamespace(const std::string& namespace_uri);
 
-/**
- * Returns an XML Android namespace for the given package of the form:
- *
- * http://schemas.android.com/apk/res/<package>
- *
- * If privateReference == true, the package will be of the form:
- *
- * http://schemas.android.com/apk/prv/res/<package>
- */
+// Returns an XML Android namespace for the given package of the form:
+//   http://schemas.android.com/apk/res/<package>
+//
+// If privateReference == true, the package will be of the form:
+//   http://schemas.android.com/apk/prv/res/<package>
 std::string BuildPackageNamespace(const android::StringPiece& package,
                                   bool private_reference = false);
 
-/**
- * Interface representing a stack of XML namespace declarations. When looking up
- * the package
- * for a namespace prefix, the stack is checked from top to bottom.
- */
+// Interface representing a stack of XML namespace declarations. When looking up the package
+// for a namespace prefix, the stack is checked from top to bottom.
 struct IPackageDeclStack {
   virtual ~IPackageDeclStack() = default;
 
-  /**
-   * Returns an ExtractedPackage struct if the alias given corresponds with a
-   * package declaration.
-   */
+  // Returns an ExtractedPackage struct if the alias given corresponds with a package declaration.
   virtual Maybe<ExtractedPackage> TransformPackageAlias(
       const android::StringPiece& alias, const android::StringPiece& local_package) const = 0;
 };
 
-/**
- * Helper function for transforming the original Reference inRef to a fully
- * qualified reference
- * via the IPackageDeclStack. This will also mark the Reference as private if
- * the namespace of the package declaration was private.
- */
+// Helper function for transforming the original Reference inRef to a fully qualified reference
+// via the IPackageDeclStack. This will also mark the Reference as private if the namespace of the
+// package declaration was private.
 void TransformReferenceFromNamespace(IPackageDeclStack* decl_stack,
                                      const android::StringPiece& local_package, Reference* in_ref);