ART: Add heap iteration callback

Add callback support for heap iteration. Visiting of fields will
be done in a follow-up.

Add a test.

Bug: 31385354
Test: m test-art-host-run-test-906-iterate-heap
Test: m ART_TEST_GC_STRESS=true ART_TEST_GC_VERIFY=true test-art-host-run-test-906-iterate-heap
Change-Id: I7bcf6751e6df4ef58756ba97701050b2ff5eb07b
diff --git a/runtime/openjdkjvmti/Android.bp b/runtime/openjdkjvmti/Android.bp
index 2d9149a..de6683c 100644
--- a/runtime/openjdkjvmti/Android.bp
+++ b/runtime/openjdkjvmti/Android.bp
@@ -18,6 +18,7 @@
     defaults: ["art_defaults"],
     host_supported: true,
     srcs: ["events.cc",
+           "heap.cc",
            "object_tagging.cc",
            "OpenjdkJvmTi.cc",
            "transform.cc"],
diff --git a/runtime/openjdkjvmti/OpenjdkJvmTi.cc b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
index 36be2a0..05da585 100644
--- a/runtime/openjdkjvmti/OpenjdkJvmTi.cc
+++ b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
@@ -39,6 +39,7 @@
 #include "art_jvmti.h"
 #include "base/mutex.h"
 #include "events-inl.h"
+#include "heap.h"
 #include "jni_env_ext-inl.h"
 #include "object_tagging.h"
 #include "obj_ptr-inl.h"
@@ -276,7 +277,8 @@
                                        jclass klass,
                                        const jvmtiHeapCallbacks* callbacks,
                                        const void* user_data) {
-    return ERR(NOT_IMPLEMENTED);
+    HeapUtil heap_util(&gObjectTagTable);
+    return heap_util.IterateThroughHeap(env, heap_filter, klass, callbacks, user_data);
   }
 
   static jvmtiError GetTag(jvmtiEnv* env, jobject object, jlong* tag_ptr) {
diff --git a/runtime/openjdkjvmti/heap.cc b/runtime/openjdkjvmti/heap.cc
new file mode 100644
index 0000000..95d9a1d
--- /dev/null
+++ b/runtime/openjdkjvmti/heap.cc
@@ -0,0 +1,166 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "heap.h"
+
+#include "art_jvmti.h"
+#include "base/macros.h"
+#include "base/mutex.h"
+#include "gc/heap.h"
+#include "mirror/class.h"
+#include "object_callbacks.h"
+#include "object_tagging.h"
+#include "obj_ptr-inl.h"
+#include "runtime.h"
+#include "scoped_thread_state_change-inl.h"
+#include "thread-inl.h"
+
+namespace openjdkjvmti {
+
+struct IterateThroughHeapData {
+  IterateThroughHeapData(HeapUtil* _heap_util,
+                         jint heap_filter,
+                         art::ObjPtr<art::mirror::Class> klass,
+                         const jvmtiHeapCallbacks* _callbacks,
+                         const void* _user_data)
+      : heap_util(_heap_util),
+        filter_klass(klass),
+        callbacks(_callbacks),
+        user_data(_user_data),
+        filter_out_tagged((heap_filter & JVMTI_HEAP_FILTER_TAGGED) != 0),
+        filter_out_untagged((heap_filter & JVMTI_HEAP_FILTER_UNTAGGED) != 0),
+        filter_out_class_tagged((heap_filter & JVMTI_HEAP_FILTER_CLASS_TAGGED) != 0),
+        filter_out_class_untagged((heap_filter & JVMTI_HEAP_FILTER_CLASS_UNTAGGED) != 0),
+        any_filter(filter_out_tagged ||
+                   filter_out_untagged ||
+                   filter_out_class_tagged ||
+                   filter_out_class_untagged),
+        stop_reports(false) {
+  }
+
+  bool ShouldReportByHeapFilter(jlong tag, jlong class_tag) {
+    if (!any_filter) {
+      return true;
+    }
+
+    if ((tag == 0 && filter_out_untagged) || (tag != 0 && filter_out_tagged)) {
+      return false;
+    }
+
+    if ((class_tag == 0 && filter_out_class_untagged) ||
+        (class_tag != 0 && filter_out_class_tagged)) {
+      return false;
+    }
+
+    return true;
+  }
+
+  HeapUtil* heap_util;
+  art::ObjPtr<art::mirror::Class> filter_klass;
+  const jvmtiHeapCallbacks* callbacks;
+  const void* user_data;
+  const bool filter_out_tagged;
+  const bool filter_out_untagged;
+  const bool filter_out_class_tagged;
+  const bool filter_out_class_untagged;
+  const bool any_filter;
+
+  bool stop_reports;
+};
+
+static void IterateThroughHeapObjectCallback(art::mirror::Object* obj, void* arg)
+    REQUIRES_SHARED(art::Locks::mutator_lock_) {
+  IterateThroughHeapData* ithd = reinterpret_cast<IterateThroughHeapData*>(arg);
+  // Early return, as we can't really stop visiting.
+  if (ithd->stop_reports) {
+    return;
+  }
+
+  art::ScopedAssertNoThreadSuspension no_suspension("IterateThroughHeapCallback");
+
+  jlong tag = 0;
+  ithd->heap_util->GetTags()->GetTag(obj, &tag);
+
+  jlong class_tag = 0;
+  art::ObjPtr<art::mirror::Class> klass = obj->GetClass();
+  ithd->heap_util->GetTags()->GetTag(klass.Ptr(), &class_tag);
+  // For simplicity, even if we find a tag = 0, assume 0 = not tagged.
+
+  if (!ithd->ShouldReportByHeapFilter(tag, class_tag)) {
+    return;
+  }
+
+  // TODO: Handle array_primitive_value_callback.
+
+  if (ithd->filter_klass != nullptr) {
+    if (ithd->filter_klass != klass) {
+      return;
+    }
+  }
+
+  jlong size = obj->SizeOf();
+
+  jint length = -1;
+  if (obj->IsArrayInstance()) {
+    length = obj->AsArray()->GetLength();
+  }
+
+  jlong saved_tag = tag;
+  jint ret = ithd->callbacks->heap_iteration_callback(class_tag,
+                                                      size,
+                                                      &tag,
+                                                      length,
+                                                      const_cast<void*>(ithd->user_data));
+
+  if (tag != saved_tag) {
+    ithd->heap_util->GetTags()->Set(obj, tag);
+  }
+
+  ithd->stop_reports = (ret & JVMTI_VISIT_ABORT) != 0;
+
+  // TODO Implement array primitive and string primitive callback.
+  // TODO Implement primitive field callback.
+}
+
+jvmtiError HeapUtil::IterateThroughHeap(jvmtiEnv* env ATTRIBUTE_UNUSED,
+                                        jint heap_filter,
+                                        jclass klass,
+                                        const jvmtiHeapCallbacks* callbacks,
+                                        const void* user_data) {
+  if (callbacks == nullptr) {
+    return ERR(NULL_POINTER);
+  }
+
+  if (callbacks->array_primitive_value_callback != nullptr) {
+    // TODO: Implement.
+    return ERR(NOT_IMPLEMENTED);
+  }
+
+  art::Thread* self = art::Thread::Current();
+  art::ScopedObjectAccess soa(self);      // Now we know we have the shared lock.
+
+  IterateThroughHeapData ithd(this,
+                              heap_filter,
+                              soa.Decode<art::mirror::Class>(klass),
+                              callbacks,
+                              user_data);
+
+  art::Runtime::Current()->GetHeap()->VisitObjects(IterateThroughHeapObjectCallback, &ithd);
+
+  return ERR(NONE);
+}
+
+}  // namespace openjdkjvmti
diff --git a/runtime/openjdkjvmti/heap.h b/runtime/openjdkjvmti/heap.h
new file mode 100644
index 0000000..fb9a216
--- /dev/null
+++ b/runtime/openjdkjvmti/heap.h
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef ART_RUNTIME_OPENJDKJVMTI_HEAP_H_
+#define ART_RUNTIME_OPENJDKJVMTI_HEAP_H_
+
+#include "jvmti.h"
+
+namespace openjdkjvmti {
+
+class ObjectTagTable;
+
+class HeapUtil {
+ public:
+  explicit HeapUtil(ObjectTagTable* tags) : tags_(tags) {
+  }
+
+  jvmtiError IterateThroughHeap(jvmtiEnv* env,
+                                jint heap_filter,
+                                jclass klass,
+                                const jvmtiHeapCallbacks* callbacks,
+                                const void* user_data);
+
+  ObjectTagTable* GetTags() {
+    return tags_;
+  }
+
+ private:
+  ObjectTagTable* tags_;
+};
+
+}  // namespace openjdkjvmti
+
+#endif  // ART_RUNTIME_OPENJDKJVMTI_HEAP_H_
diff --git a/runtime/openjdkjvmti/object_tagging.cc b/runtime/openjdkjvmti/object_tagging.cc
index bb17cac..29d4830 100644
--- a/runtime/openjdkjvmti/object_tagging.cc
+++ b/runtime/openjdkjvmti/object_tagging.cc
@@ -87,6 +87,33 @@
   return false;
 }
 
+bool ObjectTagTable::Set(art::mirror::Object* obj, jlong new_tag) {
+  art::Thread* self = art::Thread::Current();
+  art::MutexLock mu(self, allow_disallow_lock_);
+  Wait(self);
+
+  for (auto& pair : tagged_objects_) {
+    if (pair.first.Read(nullptr) == obj) {
+      pair.second = new_tag;
+      return true;
+    }
+  }
+
+  // TODO refactor with Add.
+  if (first_free_ == tagged_objects_.size()) {
+    tagged_objects_.push_back(Entry(art::GcRoot<art::mirror::Object>(obj), new_tag));
+    first_free_++;
+  } else {
+    DCHECK_LT(first_free_, tagged_objects_.size());
+    DCHECK(tagged_objects_[first_free_].first.IsNull());
+    tagged_objects_[first_free_] = Entry(art::GcRoot<art::mirror::Object>(obj), new_tag);
+    // TODO: scan for free elements.
+    first_free_ = tagged_objects_.size();
+  }
+
+  return false;
+}
+
 void ObjectTagTable::Sweep(art::IsMarkedVisitor* visitor) {
   if (event_handler_->IsEventEnabledAnywhere(JVMTI_EVENT_OBJECT_FREE)) {
     SweepImpl<true>(visitor);
diff --git a/runtime/openjdkjvmti/object_tagging.h b/runtime/openjdkjvmti/object_tagging.h
index 45f3e4f..b399e65 100644
--- a/runtime/openjdkjvmti/object_tagging.h
+++ b/runtime/openjdkjvmti/object_tagging.h
@@ -42,6 +42,10 @@
       REQUIRES_SHARED(art::Locks::mutator_lock_)
       REQUIRES(!allow_disallow_lock_);
 
+  bool Set(art::mirror::Object* obj, jlong tag)
+      REQUIRES_SHARED(art::Locks::mutator_lock_)
+      REQUIRES(!allow_disallow_lock_);
+
   bool GetTag(art::mirror::Object* obj, jlong* result)
       REQUIRES_SHARED(art::Locks::mutator_lock_)
       REQUIRES(!allow_disallow_lock_) {
diff --git a/test/906-iterate-heap/build b/test/906-iterate-heap/build
new file mode 100755
index 0000000..898e2e5
--- /dev/null
+++ b/test/906-iterate-heap/build
@@ -0,0 +1,17 @@
+#!/bin/bash
+#
+# Copyright 2016 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+./default-build "$@" --experimental agents
diff --git a/test/906-iterate-heap/expected.txt b/test/906-iterate-heap/expected.txt
new file mode 100644
index 0000000..72cd47d
--- /dev/null
+++ b/test/906-iterate-heap/expected.txt
@@ -0,0 +1,2 @@
+[{tag=1, class-tag=0, size=8, length=-1}, {tag=2, class-tag=100, size=8, length=-1}, {tag=3, class-tag=100, size=8, length=-1}, {tag=4, class-tag=0, size=32, length=5}, {tag=100, class-tag=0, size=<class>, length=-1}]
+[{tag=11, class-tag=0, size=8, length=-1}, {tag=12, class-tag=110, size=8, length=-1}, {tag=13, class-tag=110, size=8, length=-1}, {tag=14, class-tag=0, size=32, length=5}, {tag=110, class-tag=0, size=<class>, length=-1}]
diff --git a/test/906-iterate-heap/info.txt b/test/906-iterate-heap/info.txt
new file mode 100644
index 0000000..875a5f6
--- /dev/null
+++ b/test/906-iterate-heap/info.txt
@@ -0,0 +1 @@
+Tests basic functions in the jvmti plugin.
diff --git a/test/906-iterate-heap/iterate_heap.cc b/test/906-iterate-heap/iterate_heap.cc
new file mode 100644
index 0000000..ab1d8d8
--- /dev/null
+++ b/test/906-iterate-heap/iterate_heap.cc
@@ -0,0 +1,187 @@
+/*
+ * Copyright (C) 2013 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "iterate_heap.h"
+
+#include <iostream>
+#include <pthread.h>
+#include <stdio.h>
+#include <vector>
+
+#include "base/logging.h"
+#include "jni.h"
+#include "openjdkjvmti/jvmti.h"
+#include "ScopedPrimitiveArray.h"
+#include "ti-agent/common_load.h"
+
+namespace art {
+namespace Test906IterateHeap {
+
+class IterationConfig {
+ public:
+  IterationConfig() {}
+  virtual ~IterationConfig() {}
+
+  virtual jint Handle(jlong class_tag, jlong size, jlong* tag_ptr, jint length) = 0;
+};
+
+static jint JNICALL HeapIterationCallback(jlong class_tag,
+                                          jlong size,
+                                          jlong* tag_ptr,
+                                          jint length,
+                                          void* user_data) {
+  IterationConfig* config = reinterpret_cast<IterationConfig*>(user_data);
+  return config->Handle(class_tag, size, tag_ptr, length);
+}
+
+static bool Run(jint heap_filter, jclass klass_filter, IterationConfig* config) {
+  jvmtiHeapCallbacks callbacks;
+  memset(&callbacks, 0, sizeof(jvmtiHeapCallbacks));
+  callbacks.heap_iteration_callback = HeapIterationCallback;
+
+  jvmtiError ret = jvmti_env->IterateThroughHeap(heap_filter,
+                                                 klass_filter,
+                                                 &callbacks,
+                                                 config);
+  if (ret != JVMTI_ERROR_NONE) {
+    char* err;
+    jvmti_env->GetErrorName(ret, &err);
+    printf("Failure running IterateThroughHeap: %s\n", err);
+    return false;
+  }
+  return true;
+}
+
+extern "C" JNIEXPORT jint JNICALL Java_Main_iterateThroughHeapCount(JNIEnv* env ATTRIBUTE_UNUSED,
+                                                                    jclass klass ATTRIBUTE_UNUSED,
+                                                                    jint heap_filter,
+                                                                    jclass klass_filter,
+                                                                    jint stop_after) {
+  class CountIterationConfig : public IterationConfig {
+   public:
+    CountIterationConfig(jint _counter, jint _stop_after)
+        : counter(_counter),
+          stop_after(_stop_after) {
+    }
+
+    jint Handle(jlong class_tag ATTRIBUTE_UNUSED,
+                jlong size ATTRIBUTE_UNUSED,
+                jlong* tag_ptr ATTRIBUTE_UNUSED,
+                jint length ATTRIBUTE_UNUSED) OVERRIDE {
+      counter++;
+      if (counter == stop_after) {
+        return JVMTI_VISIT_ABORT;
+      }
+      return 0;
+    }
+
+    jint counter;
+    const jint stop_after;
+  };
+
+  CountIterationConfig config(0, stop_after);
+  Run(heap_filter, klass_filter, &config);
+
+  if (config.counter > config.stop_after) {
+    printf("Error: more objects visited than signaled.");
+  }
+
+  return config.counter;
+}
+
+
+extern "C" JNIEXPORT jint JNICALL Java_Main_iterateThroughHeapData(JNIEnv* env,
+                                                                   jclass klass ATTRIBUTE_UNUSED,
+                                                                   jint heap_filter,
+                                                                   jclass klass_filter,
+                                                                   jlongArray class_tags,
+                                                                   jlongArray sizes,
+                                                                   jlongArray tags,
+                                                                   jintArray lengths) {
+  class DataIterationConfig : public IterationConfig {
+   public:
+    jint Handle(jlong class_tag, jlong size, jlong* tag_ptr, jint length) OVERRIDE {
+      class_tags_.push_back(class_tag);
+      sizes_.push_back(size);
+      tags_.push_back(*tag_ptr);
+      lengths_.push_back(length);
+
+      return 0;  // Continue.
+    }
+
+    std::vector<jlong> class_tags_;
+    std::vector<jlong> sizes_;
+    std::vector<jlong> tags_;
+    std::vector<jint> lengths_;
+  };
+
+  DataIterationConfig config;
+  if (!Run(heap_filter, klass_filter, &config)) {
+    return -1;
+  }
+
+  ScopedLongArrayRW s_class_tags(env, class_tags);
+  ScopedLongArrayRW s_sizes(env, sizes);
+  ScopedLongArrayRW s_tags(env, tags);
+  ScopedIntArrayRW s_lengths(env, lengths);
+
+  for (size_t i = 0; i != config.class_tags_.size(); ++i) {
+    s_class_tags[i] = config.class_tags_[i];
+    s_sizes[i] = config.sizes_[i];
+    s_tags[i] = config.tags_[i];
+    s_lengths[i] = config.lengths_[i];
+  }
+
+  return static_cast<jint>(config.class_tags_.size());
+}
+
+extern "C" JNIEXPORT void JNICALL Java_Main_iterateThroughHeapAdd(JNIEnv* env ATTRIBUTE_UNUSED,
+                                                                  jclass klass ATTRIBUTE_UNUSED,
+                                                                  jint heap_filter,
+                                                                  jclass klass_filter) {
+  class AddIterationConfig : public IterationConfig {
+   public:
+    AddIterationConfig() {}
+
+    jint Handle(jlong class_tag ATTRIBUTE_UNUSED,
+                jlong size ATTRIBUTE_UNUSED,
+                jlong* tag_ptr,
+                jint length ATTRIBUTE_UNUSED) OVERRIDE {
+      jlong current_tag = *tag_ptr;
+      if (current_tag != 0) {
+        *tag_ptr = current_tag + 10;
+      }
+      return 0;
+    }
+  };
+
+  AddIterationConfig config;
+  Run(heap_filter, klass_filter, &config);
+}
+
+// Don't do anything
+jint OnLoad(JavaVM* vm,
+            char* options ATTRIBUTE_UNUSED,
+            void* reserved ATTRIBUTE_UNUSED) {
+  if (vm->GetEnv(reinterpret_cast<void**>(&jvmti_env), JVMTI_VERSION_1_0)) {
+    printf("Unable to get jvmti env!\n");
+    return 1;
+  }
+  return 0;
+}
+
+}  // namespace Test906IterateHeap
+}  // namespace art
diff --git a/test/906-iterate-heap/iterate_heap.h b/test/906-iterate-heap/iterate_heap.h
new file mode 100644
index 0000000..f25cdba
--- /dev/null
+++ b/test/906-iterate-heap/iterate_heap.h
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef ART_TEST_906_ITERATE_HEAP_ITERATE_HEAP_H_
+#define ART_TEST_906_ITERATE_HEAP_ITERATE_HEAP_H_
+
+#include <jni.h>
+
+namespace art {
+namespace Test906IterateHeap {
+
+jint OnLoad(JavaVM* vm, char* options, void* reserved);
+
+}  // namespace Test906IterateHeap
+}  // namespace art
+
+#endif  // ART_TEST_906_ITERATE_HEAP_ITERATE_HEAP_H_
diff --git a/test/906-iterate-heap/run b/test/906-iterate-heap/run
new file mode 100755
index 0000000..3e135a3
--- /dev/null
+++ b/test/906-iterate-heap/run
@@ -0,0 +1,43 @@
+#!/bin/bash
+#
+# Copyright 2016 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+plugin=libopenjdkjvmtid.so
+agent=libtiagentd.so
+lib=tiagentd
+if  [[ "$@" == *"-O"* ]]; then
+  agent=libtiagent.so
+  plugin=libopenjdkjvmti.so
+  lib=tiagent
+fi
+
+if [[ "$@" == *"--jvm"* ]]; then
+  arg="jvm"
+else
+  arg="art"
+fi
+
+if [[ "$@" != *"--debuggable"* ]]; then
+  other_args=" -Xcompiler-option --debuggable "
+else
+  other_args=""
+fi
+
+./default-run "$@" --experimental agents \
+                   --experimental runtime-plugins \
+                   --runtime-option -agentpath:${agent}=906-iterate-heap,${arg} \
+                   --android-runtime-option -Xplugin:${plugin} \
+                   ${other_args} \
+                   --args ${lib}
diff --git a/test/906-iterate-heap/src/Main.java b/test/906-iterate-heap/src/Main.java
new file mode 100644
index 0000000..544a365
--- /dev/null
+++ b/test/906-iterate-heap/src/Main.java
@@ -0,0 +1,146 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import java.util.ArrayList;
+import java.util.Collections;
+
+public class Main {
+  public static void main(String[] args) throws Exception {
+    System.loadLibrary(args[1]);
+
+    doTest();
+  }
+
+  public static void doTest() throws Exception {
+    A a = new A();
+    B b = new B();
+    B b2 = new B();
+    C c = new C();
+    A[] aArray = new A[5];
+
+    setTag(a, 1);
+    setTag(b, 2);
+    setTag(b2, 3);
+    setTag(aArray, 4);
+    setTag(B.class, 100);
+
+    int all = iterateThroughHeapCount(0, null, Integer.MAX_VALUE);
+    int tagged = iterateThroughHeapCount(HEAP_FILTER_OUT_UNTAGGED, null, Integer.MAX_VALUE);
+    int untagged = iterateThroughHeapCount(HEAP_FILTER_OUT_TAGGED, null, Integer.MAX_VALUE);
+    int taggedClass = iterateThroughHeapCount(HEAP_FILTER_OUT_CLASS_UNTAGGED, null,
+        Integer.MAX_VALUE);
+    int untaggedClass = iterateThroughHeapCount(HEAP_FILTER_OUT_CLASS_TAGGED, null,
+        Integer.MAX_VALUE);
+
+    if (all != tagged + untagged) {
+      throw new IllegalStateException("Instances: " + all + " != " + tagged + " + " + untagged);
+    }
+    if (all != taggedClass + untaggedClass) {
+      throw new IllegalStateException("By class: " + all + " != " + taggedClass + " + " +
+          untaggedClass);
+    }
+    if (tagged != 5) {
+      throw new IllegalStateException(tagged + " tagged objects");
+    }
+    if (taggedClass != 2) {
+      throw new IllegalStateException(tagged + " objects with tagged class");
+    }
+    if (all == tagged) {
+      throw new IllegalStateException("All objects tagged");
+    }
+    if (all == taggedClass) {
+      throw new IllegalStateException("All objects have tagged class");
+    }
+
+    long classTags[] = new long[100];
+    long sizes[] = new long[100];
+    long tags[] = new long[100];
+    int lengths[] = new int[100];
+
+    int n = iterateThroughHeapData(HEAP_FILTER_OUT_UNTAGGED, null, classTags, sizes, tags, lengths);
+    System.out.println(sort(n, classTags, sizes, tags, lengths));
+
+    iterateThroughHeapAdd(HEAP_FILTER_OUT_UNTAGGED, null);
+    n = iterateThroughHeapData(HEAP_FILTER_OUT_UNTAGGED, null, classTags, sizes, tags, lengths);
+    System.out.println(sort(n, classTags, sizes, tags, lengths));
+  }
+
+  static class A {
+  }
+
+  static class B {
+  }
+
+  static class C {
+  }
+
+  static class HeapElem implements Comparable<HeapElem> {
+    long classTag;
+    long size;
+    long tag;
+    int length;
+
+    public int compareTo(HeapElem other) {
+      if (tag != other.tag) {
+        return Long.compare(tag, other.tag);
+      }
+      if (classTag != other.classTag) {
+        return Long.compare(classTag, other.classTag);
+      }
+      if (size != other.size) {
+        return Long.compare(size, other.size);
+      }
+      return Integer.compare(length, other.length);
+    }
+
+    public String toString() {
+      return "{tag=" + tag + ", class-tag=" + classTag + ", size=" +
+          (tag >= 100 ? "<class>" : size)  // Class size is dependent on 32-bit vs 64-bit,
+                                           // so strip it.
+          + ", length=" + length + "}";
+    }
+  }
+
+  private static ArrayList<HeapElem> sort(int n, long classTags[], long sizes[], long tags[],
+      int lengths[]) {
+    ArrayList<HeapElem> ret = new ArrayList<HeapElem>(n);
+    for (int i = 0; i < n; i++) {
+      HeapElem elem = new HeapElem();
+      elem.classTag = classTags[i];
+      elem.size = sizes[i];
+      elem.tag = tags[i];
+      elem.length = lengths[i];
+      ret.add(elem);
+    }
+    Collections.sort(ret);
+    return ret;
+  }
+
+  private static native void setTag(Object o, long tag);
+  private static native long getTag(Object o);
+
+  private final static int HEAP_FILTER_OUT_TAGGED = 0x4;
+  private final static int HEAP_FILTER_OUT_UNTAGGED = 0x8;
+  private final static int HEAP_FILTER_OUT_CLASS_TAGGED = 0x10;
+  private final static int HEAP_FILTER_OUT_CLASS_UNTAGGED = 0x20;
+
+  private static native int iterateThroughHeapCount(int heapFilter,
+      Class<?> klassFilter, int stopAfter);
+  private static native int iterateThroughHeapData(int heapFilter,
+      Class<?> klassFilter, long classTags[], long sizes[], long tags[], int lengths[]);
+  private static native int iterateThroughHeapAdd(int heapFilter,
+      Class<?> klassFilter);
+}
diff --git a/test/Android.bp b/test/Android.bp
index 4457e8a..45673f5 100644
--- a/test/Android.bp
+++ b/test/Android.bp
@@ -238,8 +238,8 @@
     shared_libs: ["libartd"],
 }
 
-art_cc_test_library {
-    name: "libtiagent",
+art_cc_defaults {
+    name: "libtiagent-defaults",
     defaults: ["libartagent-defaults"],
     srcs: [
         "ti-agent/common_load.cc",
@@ -248,10 +248,18 @@
         "903-hello-tagging/tagging.cc",
         "904-object-allocation/tracking.cc",
         "905-object-free/tracking_free.cc",
+        "906-iterate-heap/iterate_heap.cc",
     ],
     shared_libs: [
-        "libart",
         "libbase",
+    ],
+}
+
+art_cc_test_library {
+    name: "libtiagent",
+    defaults: ["libtiagent-defaults"],
+    shared_libs: [
+        "libart",
         "libopenjdkjvmti",
     ],
 }
@@ -259,20 +267,11 @@
 art_cc_test_library {
     name: "libtiagentd",
     defaults: [
-        "libartagent-defaults",
+        "libtiagent-defaults",
         "art_debug_defaults",
     ],
-    srcs: [
-        "ti-agent/common_load.cc",
-        "901-hello-ti-agent/basics.cc",
-        "902-hello-transformation/transform.cc",
-        "903-hello-tagging/tagging.cc",
-        "904-object-allocation/tracking.cc",
-        "905-object-free/tracking_free.cc",
-    ],
     shared_libs: [
         "libartd",
-        "libbase",
         "libopenjdkjvmtid",
     ],
 }
diff --git a/test/ti-agent/common_load.cc b/test/ti-agent/common_load.cc
index 2b8947e..c412636 100644
--- a/test/ti-agent/common_load.cc
+++ b/test/ti-agent/common_load.cc
@@ -29,6 +29,7 @@
 #include "903-hello-tagging/tagging.h"
 #include "904-object-allocation/tracking.h"
 #include "905-object-free/tracking_free.h"
+#include "906-iterate-heap/iterate_heap.h"
 
 namespace art {
 
@@ -50,6 +51,7 @@
   { "903-hello-tagging", Test903HelloTagging::OnLoad, nullptr },
   { "904-object-allocation", Test904ObjectAllocation::OnLoad, nullptr },
   { "905-object-free", Test905ObjectFree::OnLoad, nullptr },
+  { "906-iterate-heap", Test906IterateHeap::OnLoad, nullptr },
 };
 
 static AgentLib* FindAgent(char* name) {