1 files changed, 25 insertions, 13 deletions

diff --git a/src/heap-profiler.cc b/src/heap-profiler.cc
index bfd378dd..bd1cd2d9 100644
--- a/src/heap-profiler.cc
+++ b/src/heap-profiler.cc
@@ -78,6 +78,10 @@ JSObjectsCluster Clusterizer::Clusterize(HeapObject* obj, bool fine_grain) {
     }
   } else if (obj->IsString()) {
     return JSObjectsCluster(Heap::String_symbol());
+  } else if (obj->IsJSGlobalPropertyCell()) {
+    return JSObjectsCluster(JSObjectsCluster::GLOBAL_PROPERTY);
+  } else if (obj->IsCode() || obj->IsSharedFunctionInfo() || obj->IsScript()) {
+    return JSObjectsCluster(JSObjectsCluster::CODE);
   }
   return JSObjectsCluster();
 }
@@ -112,6 +116,16 @@ int Clusterizer::CalculateNetworkSize(JSObject* obj) {
   if (FixedArray::cast(obj->elements())->length() != 0) {
     size += obj->elements()->Size();
   }
+  // For functions, also account non-empty context and literals sizes.
+  if (obj->IsJSFunction()) {
+    JSFunction* f = JSFunction::cast(obj);
+    if (f->unchecked_context()->IsContext()) {
+      size += f->context()->Size();
+    }
+    if (f->literals()->length() != 0) {
+      size += f->literals()->Size();
+    }
+  }
   return size;
 }
 
@@ -127,15 +141,15 @@ class ReferencesExtractor : public ObjectVisitor {
   }
 
   void VisitPointer(Object** o) {
-    if ((*o)->IsJSObject() || (*o)->IsString()) {
-      profile_->StoreReference(cluster_, HeapObject::cast(*o));
-    } else if ((*o)->IsFixedArray() && !inside_array_) {
+    if ((*o)->IsFixedArray() && !inside_array_) {
       // Traverse one level deep for data members that are fixed arrays.
       // This covers the case of 'elements' and 'properties' of JSObject,
       // and function contexts.
       inside_array_ = true;
       FixedArray::cast(*o)->Iterate(this);
       inside_array_ = false;
+    } else if ((*o)->IsHeapObject()) {
+      profile_->StoreReference(cluster_, HeapObject::cast(*o));
     }
   }
 
@@ -340,6 +354,8 @@ void JSObjectsCluster::Print(StringStream* accumulator) const {
     accumulator->Add("(roots)");
   } else if (constructor_ == FromSpecialCase(GLOBAL_PROPERTY)) {
     accumulator->Add("(global property)");
+  } else if (constructor_ == FromSpecialCase(CODE)) {
+    accumulator->Add("(code)");
   } else if (constructor_ == FromSpecialCase(SELF)) {
     accumulator->Add("(self)");
   } else {
@@ -520,13 +536,14 @@ RetainerHeapProfile::RetainerHeapProfile()
     : zscope_(DELETE_ON_EXIT) {
   JSObjectsCluster roots(JSObjectsCluster::ROOTS);
   ReferencesExtractor extractor(roots, this);
-  Heap::IterateRoots(&extractor);
+  Heap::IterateRoots(&extractor, VISIT_ONLY_STRONG);
 }
 
 
 void RetainerHeapProfile::StoreReference(const JSObjectsCluster& cluster,
                                          HeapObject* ref) {
   JSObjectsCluster ref_cluster = Clusterizer::Clusterize(ref);
+  if (ref_cluster.is_null()) return;
   JSObjectsRetainerTree::Locator ref_loc;
   if (retainers_tree_.Insert(ref_cluster, &ref_loc)) {
     ref_loc.set_value(new JSObjectsClusterTree());
@@ -537,15 +554,10 @@ void RetainerHeapProfile::StoreReference(const JSObjectsCluster& cluster,
 
 
 void RetainerHeapProfile::CollectStats(HeapObject* obj) {
-  if (obj->IsJSObject()) {
-    const JSObjectsCluster cluster = Clusterizer::Clusterize(obj);
-    ReferencesExtractor extractor(cluster, this);
-    obj->Iterate(&extractor);
-  } else if (obj->IsJSGlobalPropertyCell()) {
-    JSObjectsCluster global_prop(JSObjectsCluster::GLOBAL_PROPERTY);
-    ReferencesExtractor extractor(global_prop, this);
-    obj->Iterate(&extractor);
-  }
+  const JSObjectsCluster cluster = Clusterizer::Clusterize(obj);
+  if (cluster.is_null()) return;
+  ReferencesExtractor extractor(cluster, this);
+  obj->Iterate(&extractor);
 }