Bug 1215397 - Add state and UI for breakdowns in memory tool. r=fitzgen

devtools/client/memory/actions/breakdown.js

@@ -0,0 +1,43 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at http://mozilla.org/MPL/2.0/. */
+"use strict";
+
+// @TODO 1215606
+// Use this assert instead of utils when fixed.
+// const { assert } = require("devtools/shared/DevToolsUtils");
+const { breakdownEquals, createSnapshot, assert } = require("../utils");
+const { actions, snapshotState: states } = require("../constants");
+const { takeCensus } = require("./snapshot");
+
+const setBreakdownAndRefresh = exports.setBreakdownAndRefresh = function (heapWorker, breakdown) {
+  return function *(dispatch, getState) {
+    // Clears out all stored census data and sets
+    // the breakdown
+    dispatch(setBreakdown(breakdown));
+    let snapshot = getState().snapshots.find(s => s.selected);
+
+    // If selected snapshot does not have updated census if the breakdown
+    // changed, retake the census with new breakdown
+    if (snapshot && !breakdownEquals(snapshot.breakdown, breakdown)) {
+      yield dispatch(takeCensus(heapWorker, snapshot));
+    }
+  };
+};
+
+/**
+ * Clears out all census data in the snapshots and sets
+ * a new breakdown.
+ *
+ * @param {Breakdown} breakdown
+ */
+const setBreakdown = exports.setBreakdown = function (breakdown) {
+  // @TODO 1215606
+  assert(typeof breakdown === "object" && breakdown.by,
+    `Breakdowns must be an object with a \`by\` property, attempted to set: ${uneval(breakdown)}`);
+
+  return {
+    type: actions.SET_BREAKDOWN,
+    breakdown,
+  }
+};

devtools/client/memory/actions/moz.build

@@ -4,5 +4,6 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 DevToolsModules(
+    'breakdown.js',
     'snapshot.js',
 )

devtools/client/memory/actions/snapshot.js

@@ -6,7 +6,7 @@
 // @TODO 1215606
 // Use this assert instead of utils when fixed.
 // const { assert } = require("devtools/shared/DevToolsUtils");
-const { createSnapshot, assert } = require("../utils");
+const { getSnapshot, breakdownEquals, createSnapshot, assert } = require("../utils");
 const { actions, snapshotState: states } = require("../constants");
 
 /**
@@ -17,19 +17,36 @@ const { actions, snapshotState: states } = require("../constants");
  * @param {HeapAnalysesClient}
  * @param {Object}
  */
-const takeSnapshotAndCensus = exports.takeSnapshotAndCensus = function takeSnapshotAndCensus (front, heapWorker) {
-  return function *(dispatch, getStore) {
+const takeSnapshotAndCensus = exports.takeSnapshotAndCensus = function (front, heapWorker) {
+  return function *(dispatch, getState) {
     let snapshot = yield dispatch(takeSnapshot(front));
     yield dispatch(readSnapshot(heapWorker, snapshot));
     yield dispatch(takeCensus(heapWorker, snapshot));
   };
 };
 
+/**
+ * Selects a snapshot and if the snapshot's census is using a different
+ * breakdown, take a new census.
+ *
+ * @param {HeapAnalysesClient}
+ * @param {Snapshot}
+ */
+const selectSnapshotAndRefresh = exports.selectSnapshotAndRefresh = function (heapWorker, snapshot) {
+  return function *(dispatch, getState) {
+    dispatch(selectSnapshot(snapshot));
+
+    // Attempt to take another census; if the snapshot already is using
+    // the correct breakdown, this will noop.
+    yield dispatch(takeCensus(heapWorker, snapshot));
+  };
+};
+
 /**
  * @param {MemoryFront}
  */
-const takeSnapshot = exports.takeSnapshot = function takeSnapshot (front) {
-  return function *(dispatch, getStore) {
+const takeSnapshot = exports.takeSnapshot = function (front) {
+  return function *(dispatch, getState) {
     let snapshot = createSnapshot();
     dispatch({ type: actions.TAKE_SNAPSHOT_START, snapshot });
     dispatch(selectSnapshot(snapshot));
@@ -49,7 +66,7 @@ const takeSnapshot = exports.takeSnapshot = function takeSnapshot (front) {
  * @param {Snapshot} snapshot,
  */
 const readSnapshot = exports.readSnapshot = function readSnapshot (heapWorker, snapshot) {
-  return function *(dispatch, getStore) {
+  return function *(dispatch, getState) {
     // @TODO 1215606
     assert(snapshot.state === states.SAVED,
       "Should only read a snapshot once");
@@ -64,29 +81,42 @@ const readSnapshot = exports.readSnapshot = function readSnapshot (heapWorker, s
  * @param {HeapAnalysesClient} heapWorker
  * @param {Snapshot} snapshot,
  *
- * @see {Snapshot} model defined in devtools/client/memory/app.js
+ * @see {Snapshot} model defined in devtools/client/memory/models.js
  * @see `devtools/shared/heapsnapshot/HeapAnalysesClient.js`
  * @see `js/src/doc/Debugger/Debugger.Memory.md` for breakdown details
  */
-const takeCensus = exports.takeCensus = function takeCensus (heapWorker, snapshot) {
-  return function *(dispatch, getStore) {
+const takeCensus = exports.takeCensus = function (heapWorker, snapshot) {
+  return function *(dispatch, getState) {
     // @TODO 1215606
     assert([states.READ, states.SAVED_CENSUS].includes(snapshot.state),
       "Can only take census of snapshots in READ or SAVED_CENSUS state");
 
-    let breakdown = getStore().breakdown;
-    dispatch({ type: actions.TAKE_CENSUS_START, snapshot, breakdown });
+    let census;
+    let breakdown = getState().breakdown;
 
-    let census = yield heapWorker.takeCensus(snapshot.path, { breakdown }, { asTreeNode: true });
-    dispatch({ type: actions.TAKE_CENSUS_END, snapshot, census });
+    // If breakdown hasn't changed, don't do anything
+    if (breakdownEquals(breakdown, snapshot.breakdown)) {
+      return;
+    }
+
+    // Keep taking a census if the breakdown changes during. Recheck
+    // that the breakdown used for the census is the same as
+    // the state's breakdown.
+    do {
+      breakdown = getState().breakdown;
+      dispatch({ type: actions.TAKE_CENSUS_START, snapshot, breakdown });
+      census = yield heapWorker.takeCensus(snapshot.path, { breakdown }, { asTreeNode: true });
+    } while (!breakdownEquals(breakdown, getState().breakdown));
+
+    dispatch({ type: actions.TAKE_CENSUS_END, snapshot, breakdown, census });
   };
 };
 
 /**
  * @param {Snapshot}
- * @see {Snapshot} model defined in devtools/client/memory/app.js
+ * @see {Snapshot} model defined in devtools/client/memory/models.js
  */
-const selectSnapshot = exports.selectSnapshot = function takeSnapshot (snapshot) {
+const selectSnapshot = exports.selectSnapshot = function (snapshot) {
   return {
     type: actions.SELECT_SNAPSHOT,
     snapshot

devtools/client/memory/app.js

@@ -1,59 +1,18 @@
 const { DOM: dom, createClass, createFactory, PropTypes } = require("devtools/client/shared/vendor/react");
 const { connect } = require("devtools/client/shared/vendor/react-redux");
-const { selectSnapshot, takeSnapshotAndCensus } = require("./actions/snapshot");
-const { snapshotState } = require("./constants");
+const { selectSnapshotAndRefresh, takeSnapshotAndCensus } = require("./actions/snapshot");
+const { setBreakdownAndRefresh } = require("./actions/breakdown");
+const { breakdownNameToSpec, getBreakdownDisplayData } = require("./utils");
 const Toolbar = createFactory(require("./components/toolbar"));
 const List = createFactory(require("./components/list"));
 const SnapshotListItem = createFactory(require("./components/snapshot-list-item"));
 const HeapView = createFactory(require("./components/heap"));
-
-const stateModel = {
-  /**
-   * {MemoryFront}
-   * Used to communicate with the platform.
-   */
-  front: PropTypes.any,
-
-  /**
-   * {HeapAnalysesClient}
-   * Used to communicate with the worker that performs analyses on heaps.
-   */
-  heapWorker: PropTypes.any,
-
-  /**
-   * The breakdown object DSL describing how we want
-   * the census data to be.
-   * @see `js/src/doc/Debugger/Debugger.Memory.md`
-   */
-  breakdown: PropTypes.object.isRequired,
-
-  /**
-   * {Array<Snapshot>}
-   * List of references to all snapshots taken
-   */
-  snapshots: PropTypes.arrayOf(PropTypes.shape({
-    // Unique ID for a snapshot
-    id: PropTypes.number.isRequired,
-    // fs path to where the snapshot is stored; used to
-    // identify the snapshot for HeapAnalysesClient.
-    path: PropTypes.string,
-    // Whether or not this snapshot is currently selected.
-    selected: PropTypes.bool.isRequired,
-    // Whther or not the snapshot has been read into memory.
-    // Only needed to do once.
-    snapshotRead: PropTypes.bool.isRequired,
-    // State the snapshot is in
-    // @see ./constants.js
-    state: PropTypes.oneOf(Object.keys(snapshotState)).isRequired,
-    // Data of a census breakdown
-    census: PropTypes.any,
-  }))
-};
+const { app: appModel } = require("./models");
 
 const App = createClass({
   displayName: "memory-tool",
 
-  propTypes: stateModel,
+  propTypes: appModel,
 
   childContextTypes: {
     front: PropTypes.any,
@@ -75,17 +34,17 @@ const App = createClass({
       dom.div({ id: "memory-tool" }, [
 
         Toolbar({
-          buttons: [{
-            className: "take-snapshot",
-            onClick: () => dispatch(takeSnapshotAndCensus(front, heapWorker))
-          }]
+          breakdowns: getBreakdownDisplayData(),
+          onTakeSnapshotClick: () => dispatch(takeSnapshotAndCensus(front, heapWorker)),
+          onBreakdownChange: breakdown =>
+            dispatch(setBreakdownAndRefresh(heapWorker, breakdownNameToSpec(breakdown))),
         }),
 
         dom.div({ id: "memory-tool-container" }, [
           List({
             itemComponent: SnapshotListItem,
             items: snapshots,
-            onClick: snapshot => dispatch(selectSnapshot(snapshot))
+            onClick: snapshot => dispatch(selectSnapshotAndRefresh(heapWorker, snapshot))
           }),
 
           HeapView({

devtools/client/memory/components/heap.js

@@ -1,6 +1,7 @@
 const { DOM: dom, createClass, PropTypes } = require("devtools/client/shared/vendor/react");
 const { getSnapshotStatusText } = require("../utils");
 const { snapshotState: states } = require("../constants");
+const { snapshot: snapshotModel } = require("../models");
 const TAKE_SNAPSHOT_TEXT = "Take snapshot";
 
 /**
@@ -14,7 +15,7 @@ const Heap = module.exports = createClass({
 
   propTypes: {
     onSnapshotClick: PropTypes.func.isRequired,
-    snapshot: PropTypes.any,
+    snapshot: snapshotModel,
   },
 
   render() {
@@ -22,7 +23,6 @@ const Heap = module.exports = createClass({
     let pane;
     let census = snapshot ? snapshot.census : null;
     let state = snapshot ? snapshot.state : "initial";
-    let statusText = getSnapshotStatusText(snapshot);
 
     switch (state) {
       case "initial":
@@ -35,7 +35,8 @@ const Heap = module.exports = createClass({
       case states.READING:
       case states.READ:
       case states.SAVING_CENSUS:
-        pane = dom.div({ className: "heap-view-panel", "data-state": state }, statusText);
+        pane = dom.div({ className: "heap-view-panel", "data-state": state },
+          getSnapshotStatusText(snapshot));
         break;
       case states.SAVED_CENSUS:
         pane = dom.div({ className: "heap-view-panel", "data-state": "loaded" }, JSON.stringify(census || {}));

devtools/client/memory/components/snapshot-list-item.js

@@ -1,12 +1,13 @@
 const { DOM: dom, createClass, PropTypes } = require("devtools/client/shared/vendor/react");
 const { getSnapshotStatusText } = require("../utils");
+const { snapshot: snapshotModel } = require("../models");
 
 const SnapshotListItem = module.exports = createClass({
   displayName: "snapshot-list-item",
 
   propTypes: {
     onClick: PropTypes.func,
-    item: PropTypes.any.isRequired,
+    item: snapshotModel.isRequired,
     index: PropTypes.number.isRequired,
   },
 

devtools/client/memory/components/toolbar.js

@@ -1,16 +1,26 @@
-const { DOM, createClass } = require("devtools/client/shared/vendor/react");
+const { DOM, createClass, PropTypes } = require("devtools/client/shared/vendor/react");
 
 const Toolbar = module.exports = createClass({
   displayName: "toolbar",
+  propTypes: {
+    breakdowns: PropTypes.arrayOf(PropTypes.shape({
+      name: PropTypes.string.isRequired,
+      displayName: PropTypes.string.isRequired,
+    })).isRequired,
+    onTakeSnapshotClick: PropTypes.func.isRequired,
+    onBreakdownChange: PropTypes.func.isRequired,
+  },
 
   render() {
-    let buttons = this.props.buttons;
+    let { onTakeSnapshotClick, onBreakdownChange, breakdowns } = this.props;
     return (
-      DOM.div({ className: "devtools-toolbar" }, ...buttons.map(spec => {
-        return DOM.button(Object.assign({}, spec, {
-          className: `${spec.className || "" } devtools-button`
-        }));
-      }))
+      DOM.div({ className: "devtools-toolbar" }, [
+        DOM.button({ className: `take-snapshot devtools-button`, onClick: onTakeSnapshotClick }),
+        DOM.select({
+          className: `select-breakdown`,
+          onChange: e => onBreakdownChange(e.target.value),
+        }, breakdowns.map(({ name, displayName }) => DOM.option({ value: name }, displayName)))
+      ])
     );
   }
 });

devtools/client/memory/constants.js

@@ -21,6 +21,39 @@ actions.TAKE_CENSUS_END = "take-census-end";
 // Fired by UI to select a snapshot to view.
 actions.SELECT_SNAPSHOT = "select-snapshot";
 
+const COUNT = { by: "count", count: true, bytes: true };
+const INTERNAL_TYPE = { by: "internalType", then: COUNT };
+const ALLOCATION_STACK = { by: "allocationStack", then: COUNT, noStack: COUNT };
+const OBJECT_CLASS = { by: "objectClass", then: COUNT, other: COUNT };
+
+const breakdowns = exports.breakdowns = {
+  coarseType: {
+    displayName: "Coarse Type",
+    breakdown: {
+      by: "coarseType",
+      objects: ALLOCATION_STACK,
+      strings: ALLOCATION_STACK,
+      scripts: INTERNAL_TYPE,
+      other: INTERNAL_TYPE,
+    }
+  },
+
+  allocationStack: {
+    displayName: "Allocation Site",
+    breakdown: ALLOCATION_STACK,
+  },
+
+  objectClass: {
+    displayName: "Object Class",
+    breakdown: OBJECT_CLASS,
+  },
+
+  internalType: {
+    displayName: "Internal Type",
+    breakdown: INTERNAL_TYPE,
+  },
+};
+
 const snapshotState = exports.snapshotState = {};
 
 /**

devtools/client/memory/models.js

@@ -0,0 +1,61 @@
+const { MemoryFront } = require("devtools/server/actors/memory");
+const HeapAnalysesClient = require("devtools/shared/heapsnapshot/HeapAnalysesClient");
+const { PropTypes } = require("devtools/client/shared/vendor/react");
+const { snapshotState: states } = require("./constants");
+
+/**
+ * The breakdown object DSL describing how we want
+ * the census data to be.
+ * @see `js/src/doc/Debugger/Debugger.Memory.md`
+ */
+let breakdownModel = exports.breakdown = PropTypes.shape({
+  by: PropTypes.oneOf(["coarseType", "allocationStack", "objectClass", "internalType"]).isRequired,
+});
+
+/**
+ * Snapshot model.
+ */
+let snapshotModel = exports.snapshot = PropTypes.shape({
+  // Unique ID for a snapshot
+  id: PropTypes.number.isRequired,
+  // Whether or not this snapshot is currently selected.
+  selected: PropTypes.bool.isRequired,
+  // fs path to where the snapshot is stored; used to
+  // identify the snapshot for HeapAnalysesClient.
+  path: PropTypes.string,
+  // Data of a census breakdown
+  census: PropTypes.object,
+  // The breakdown used to generate the current census
+  breakdown: breakdownModel,
+  // State the snapshot is in
+  // @see ./constants.js
+  state: function (props, propName) {
+    let stateNames = Object.keys(states);
+    let current = props.state;
+    let shouldHavePath = [states.SAVED, states.READ, states.SAVING_CENSUS, states.SAVED_CENSUS];
+    let shouldHaveCensus = [states.SAVED_CENSUS];
+
+    if (!stateNames.contains(current)) {
+      throw new Error(`Snapshot state must be one of ${stateNames}.`);
+    }
+    if (shouldHavePath.contains(current) && !path) {
+      throw new Error(`Snapshots in state ${current} must have a snapshot path.`);
+    }
+    if (shouldHaveCensus.contains(current) && (!props.census || !props.breakdown)) {
+      throw new Error(`Snapshots in state ${current} must have a census and breakdown.`);
+    }
+  },
+});
+
+let appModel = exports.app = {
+  // {MemoryFront} Used to communicate with platform
+  front: PropTypes.instanceOf(MemoryFront),
+  // {HeapAnalysesClient} Used to interface with snapshots
+  heapWorker: PropTypes.instanceOf(HeapAnalysesClient),
+  // The breakdown object DSL describing how we want
+  // the census data to be.
+  // @see `js/src/doc/Debugger/Debugger.Memory.md`
+  breakdown: breakdownModel.isRequired,
+  // List of reference to all snapshots taken
+  snapshots: PropTypes.arrayOf(snapshotModel).isRequired,
+};

devtools/client/memory/moz.build

@@ -14,6 +14,7 @@ DevToolsModules(
     'app.js',
     'constants.js',
     'initializer.js',
+    'models.js',
     'panel.js',
     'reducers.js',
     'store.js',

devtools/client/memory/reducers/breakdown.js

@@ -1,15 +1,16 @@
-const { actions } = require("../constants");
+const { actions, breakdowns } = require("../constants");
+const DEFAULT_BREAKDOWN = breakdowns.coarseType.breakdown;
 
-// Hardcoded breakdown for now
-const DEFAULT_BREAKDOWN = {
-  by: "internalType",
-  then: { by: "count", count: true, bytes: true }
+let handlers = Object.create(null);
+
+handlers[actions.SET_BREAKDOWN] = function (_, action) {
+  return Object.assign({}, action.breakdown);
 };
 
-/**
- * Not much to do here yet until we can change breakdowns,
- * but this gets it in our store.
- */
 module.exports = function (state=DEFAULT_BREAKDOWN, action) {
-  return Object.assign({}, DEFAULT_BREAKDOWN);
+  let handle = handlers[action.type];
+  if (handle) {
+    return handle(state, action);
+  }
+  return state;
 };

devtools/client/memory/reducers/snapshots.js

@@ -33,6 +33,7 @@ handlers[actions.TAKE_CENSUS_START] = function (snapshots, action) {
   let snapshot = getSnapshot(snapshots, action.snapshot);
   snapshot.state = states.SAVING_CENSUS;
   snapshot.census = null;
+  snapshot.breakdown = action.breakdown;
   return [...snapshots];
 };
 
@@ -40,6 +41,7 @@ handlers[actions.TAKE_CENSUS_END] = function (snapshots, action) {
   let snapshot = getSnapshot(snapshots, action.snapshot);
   snapshot.state = states.SAVED_CENSUS;
   snapshot.census = action.census;
+  snapshot.breakdown = action.breakdown;
   return [...snapshots];
 };
 

devtools/client/memory/test/unit/head.js

@@ -59,3 +59,32 @@ function waitUntilState (store, predicate) {
 
   return deferred.promise;
 }
+
+function waitUntilSnapshotState (store, expected) {
+  let predicate = () => {
+    let snapshots = store.getState().snapshots;
+    do_print(snapshots.map(x => x.state));
+    return snapshots.length === expected.length &&
+           expected.every((state, i) => state === "*" || snapshots[i].state === state);
+  };
+  do_print(`Waiting for snapshots to be of state: ${expected}`);
+  return waitUntilState(store, predicate);
+}
+
+function isBreakdownType (census, type) {
+  // Little sanity check, all censuses should have atleast a children array
+  if (!census || !Array.isArray(census.children)) {
+    return false;
+  }
+  switch (type) {
+    case "coarseType":
+      return census.children.find(c => c.name === "objects");
+    case "objectClass":
+      return census.children.find(c => c.name === "Function");
+    case "internalType":
+      return census.children.find(c => c.name === "js::BaseShape") &&
+             !census.children.find(c => c.name === "objects");
+    default:
+      throw new Error(`isBreakdownType does not yet support ${type}`);
+  }
+}

devtools/client/memory/test/unit/test_action-set-breakdown-and-refresh-01.js

@@ -0,0 +1,92 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+/**
+ * Tests the task creator `setBreakdownAndRefreshAndRefresh()` for breakdown changing.
+ * We test this rather than `setBreakdownAndRefresh` directly, as we use the refresh action
+ * in the app itself composed from `setBreakdownAndRefresh`
+ */
+
+let { breakdowns, snapshotState: states } = require("devtools/client/memory/constants");
+let { breakdownEquals } = require("devtools/client/memory/utils");
+let { setBreakdownAndRefresh } = require("devtools/client/memory/actions/breakdown");
+let { takeSnapshotAndCensus, selectSnapshotAndRefresh } = require("devtools/client/memory/actions/snapshot");
+
+function run_test() {
+  run_next_test();
+}
+
+add_task(function *() {
+  let front = new StubbedMemoryFront();
+  let heapWorker = new HeapAnalysesClient();
+  yield front.attach();
+  let store = Store();
+  let { getState, dispatch } = store;
+
+  // Test default breakdown with no snapshots
+  equal(getState().breakdown.by, "coarseType", "default coarseType breakdown selected at start.");
+  dispatch(setBreakdownAndRefresh(heapWorker, breakdowns.objectClass.breakdown));
+  equal(getState().breakdown.by, "objectClass", "breakdown changed with no snapshots");
+
+  // Test invalid breakdowns
+  ok(getState().errors.length === 0, "No error actions in the queue.");
+  dispatch(setBreakdownAndRefresh(heapWorker, {}));
+  yield waitUntilState(store, () => getState().errors.length === 1);
+  ok(true, "Emits an error action when passing in an invalid breakdown object");
+
+  equal(getState().breakdown.by, "objectClass",
+    "current breakdown unchanged when passing invalid breakdown");
+
+  // Test new snapshots
+  dispatch(takeSnapshotAndCensus(front, heapWorker));
+  yield waitUntilSnapshotState(store, [states.SAVED_CENSUS]);
+  ok(isBreakdownType(getState().snapshots[0].census, "objectClass"),
+    "New snapshots use the current, non-default breakdown");
+
+
+  // Updates when changing breakdown during `SAVING`
+  dispatch(takeSnapshotAndCensus(front, heapWorker));
+  yield waitUntilSnapshotState(store, [states.SAVED_CENSUS, states.SAVING]);
+  dispatch(setBreakdownAndRefresh(heapWorker, breakdowns.coarseType.breakdown));
+  yield waitUntilSnapshotState(store, [states.SAVED_CENSUS, states.SAVED_CENSUS]);
+
+  ok(isBreakdownType(getState().snapshots[1].census, "coarseType"),
+    "Breakdown can be changed while saving snapshots, uses updated breakdown in census");
+
+
+  // Updates when changing breakdown during `SAVING_CENSUS`
+  dispatch(takeSnapshotAndCensus(front, heapWorker));
+  yield waitUntilSnapshotState(store, [states.SAVED_CENSUS, states.SAVED_CENSUS, states.SAVING_CENSUS]);
+  dispatch(setBreakdownAndRefresh(heapWorker, breakdowns.objectClass.breakdown));
+  yield waitUntilSnapshotState(store, [states.SAVED_CENSUS, states.SAVED_CENSUS, states.SAVED_CENSUS]);
+
+  ok(breakdownEquals(getState().snapshots[2].breakdown, breakdowns.objectClass.breakdown),
+    "Breakdown can be changed while saving census, stores updated breakdown in snapshot");
+  ok(isBreakdownType(getState().snapshots[2].census, "objectClass"),
+    "Breakdown can be changed while saving census, uses updated breakdown in census");
+
+  // Updates census on currently selected snapshot when changing breakdown
+  ok(getState().snapshots[2].selected, "Third snapshot currently selected");
+  dispatch(setBreakdownAndRefresh(heapWorker, breakdowns.internalType.breakdown));
+  yield waitUntilState(store, () => isBreakdownType(getState().snapshots[2].census, "internalType"));
+  ok(isBreakdownType(getState().snapshots[2].census, "internalType"),
+    "Snapshot census updated when changing breakdowns after already generating one census");
+
+  // Does not update unselected censuses
+  ok(!getState().snapshots[1].selected, "Second snapshot unselected currently");
+  ok(breakdownEquals(getState().snapshots[1].breakdown, breakdowns.coarseType.breakdown),
+    "Second snapshot using `coarseType` breakdown still and not yet updated to correct breakdown");
+  ok(isBreakdownType(getState().snapshots[1].census, "coarseType"),
+    "Second snapshot using `coarseType` still for census and not yet updated to correct breakdown");
+
+  // Updates to current breakdown when switching to stale snapshot
+  dispatch(selectSnapshotAndRefresh(heapWorker, getState().snapshots[1]));
+  yield waitUntilSnapshotState(store, [states.SAVED_CENSUS, states.SAVING_CENSUS, states.SAVED_CENSUS]);
+  yield waitUntilSnapshotState(store, [states.SAVED_CENSUS, states.SAVED_CENSUS, states.SAVED_CENSUS]);
+
+  ok(getState().snapshots[1].selected, "Second snapshot selected currently");
+  ok(breakdownEquals(getState().snapshots[1].breakdown, breakdowns.internalType.breakdown),
+    "Second snapshot using `internalType` breakdown and updated to correct breakdown");
+  ok(isBreakdownType(getState().snapshots[1].census, "internalType"),
+    "Second snapshot using `internalType` for census and updated to correct breakdown");
+});

devtools/client/memory/test/unit/test_action-set-breakdown-and-refresh-02.js

@@ -0,0 +1,38 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+/**
+ * Tests the task creator `setBreakdownAndRefreshAndRefresh()` for custom
+ * breakdowns.
+ */
+
+let { snapshotState: states } = require("devtools/client/memory/constants");
+let { breakdownEquals } = require("devtools/client/memory/utils");
+let { setBreakdownAndRefresh } = require("devtools/client/memory/actions/breakdown");
+let { takeSnapshotAndCensus } = require("devtools/client/memory/actions/snapshot");
+let custom = { by: "internalType", then: { by: "count", bytes: true }};
+
+function run_test() {
+  run_next_test();
+}
+
+add_task(function *() {
+  let front = new StubbedMemoryFront();
+  let heapWorker = new HeapAnalysesClient();
+  yield front.attach();
+  let store = Store();
+  let { getState, dispatch } = store;
+
+  dispatch(setBreakdownAndRefresh(heapWorker, custom));
+  ok(breakdownEquals(getState().breakdown, custom),
+    "Custom breakdown stored in breakdown state.");
+
+  dispatch(takeSnapshotAndCensus(front, heapWorker));
+  yield waitUntilSnapshotState(store, [states.SAVED_CENSUS]);
+
+  ok(breakdownEquals(getState().snapshots[0].breakdown, custom),
+    "New snapshot stored custom breakdown when done taking census");
+  ok(getState().snapshots[0].census.children.length, "Census has some children");
+  // Ensure we don't have `count` in any results
+  ok(getState().snapshots[0].census.children.every(c => !c.count), "Census used custom breakdown");
+});

devtools/client/memory/test/unit/test_action-set-breakdown.js

@@ -0,0 +1,45 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+/**
+ * Tests the action creator `setBreakdown()` for breakdown changing.
+ * Does not test refreshing the census information, check `setBreakdownAndRefresh` action
+ * for that.
+ */
+
+let { breakdowns, snapshotState: states } = require("devtools/client/memory/constants");
+let { setBreakdown } = require("devtools/client/memory/actions/breakdown");
+let { takeSnapshotAndCensus } = require("devtools/client/memory/actions/snapshot");
+
+function run_test() {
+  run_next_test();
+}
+
+add_task(function *() {
+  let front = new StubbedMemoryFront();
+  let heapWorker = new HeapAnalysesClient();
+  yield front.attach();
+  let store = Store();
+  let { getState, dispatch } = store;
+
+  // Test default breakdown with no snapshots
+  equal(getState().breakdown.by, "coarseType", "default coarseType breakdown selected at start.");
+  dispatch(setBreakdown(breakdowns.objectClass.breakdown));
+  equal(getState().breakdown.by, "objectClass", "breakdown changed with no snapshots");
+
+  // Test invalid breakdowns
+  try {
+    dispatch(setBreakdown({}));
+    ok(false, "Throws when passing in an invalid breakdown object");
+  } catch (e) {
+    ok(true, "Throws when passing in an invalid breakdown object");
+  }
+  equal(getState().breakdown.by, "objectClass",
+    "current breakdown unchanged when passing invalid breakdown");
+
+  // Test new snapshots
+  dispatch(takeSnapshotAndCensus(front, heapWorker));
+  yield waitUntilSnapshotState(store, [states.SAVED_CENSUS]);
+  ok(isBreakdownType(getState().snapshots[0].census, "objectClass"),
+    "New snapshots use the current, non-default breakdown");
+});

devtools/client/memory/test/unit/test_action-take-census.js

@@ -5,7 +5,8 @@
  * Tests the async reducer responding to the action `takeCensus(heapWorker, snapshot)`
  */
 
-var { snapshotState: states } = require("devtools/client/memory/constants");
+var { snapshotState: states, breakdowns } = require("devtools/client/memory/constants");
+var { breakdownEquals } = require("devtools/client/memory/utils");
 var { ERROR_TYPE } = require("devtools/client/shared/redux/middleware/task");
 var actions = require("devtools/client/memory/actions/snapshot");
 
@@ -43,7 +44,9 @@ add_task(function *() {
 
   snapshot = store.getState().snapshots[0];
   ok(snapshot.census, "Snapshot has census after saved census");
-  ok(snapshot.census.children.length, "Census is in tree node form with the default breakdown");
-  ok(snapshot.census.children.find(t => t.name === "JSObject"),
+  ok(snapshot.census.children.length, "Census is in tree node form");
+  ok(isBreakdownType(snapshot.census, "coarseType"),
     "Census is in tree node form with the default breakdown");
+  ok(breakdownEquals(snapshot.breakdown, breakdowns.coarseType.breakdown),
+    "Snapshot stored correct breakdown used for the census");
 });

devtools/client/memory/test/unit/test_utils.js

@@ -0,0 +1,51 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+/**
+ * Tests the task creator `takeSnapshotAndCensus()` for the whole flow of
+ * taking a snapshot, and its sub-actions.
+ */
+
+let utils = require("devtools/client/memory/utils");
+let { snapshotState: states, breakdowns } = require("devtools/client/memory/constants");
+let { Preferences } = require("resource://gre/modules/Preferences.jsm");
+
+function run_test() {
+  run_next_test();
+}
+
+add_task(function *() {
+  ok(utils.breakdownEquals(breakdowns.allocationStack.breakdown, {
+    by: "allocationStack",
+    then: { by: "count", count: true, bytes: true },
+    noStack: { by: "count", count: true, bytes: true },
+  }), "utils.breakdownEquals() passes with preset"),
+
+  ok(!utils.breakdownEquals(breakdowns.allocationStack.breakdown, {
+    by: "allocationStack",
+    then: { by: "count", count: false, bytes: true },
+    noStack: { by: "count", count: true, bytes: true },
+  }), "utils.breakdownEquals() fails when deep properties do not match");
+
+  ok(!utils.breakdownEquals(breakdowns.allocationStack.breakdown, {
+    by: "allocationStack",
+    then: { by: "count", bytes: true },
+    noStack: { by: "count", count: true, bytes: true },
+  }), "utils.breakdownEquals() fails when deep properties are missing.");
+
+  let s1 = utils.createSnapshot();
+  let s2 = utils.createSnapshot();
+  ok(s1.state, states.SAVING, "utils.createSnapshot() creates snapshot in saving state");
+  ok(s1.id !== s2.id, "utils.createSnapshot() creates snapshot with unique ids");
+
+  ok(utils.breakdownEquals(utils.breakdownNameToSpec("coarseType"), breakdowns.coarseType.breakdown),
+    "utils.breakdownNameToSpec() works for presets");
+  ok(utils.breakdownEquals(utils.breakdownNameToSpec("coarseType"), breakdowns.coarseType.breakdown),
+    "utils.breakdownNameToSpec() works for presets");
+
+  let custom = { by: "internalType", then: { by: "count", bytes: true }};
+  Preferences.set("devtools.memory.custom-breakdowns", JSON.stringify({ "My Breakdown": custom }));
+
+  ok(utils.breakdownEquals(utils.getCustomBreakdowns()["My Breakdown"], custom),
+    "utils.getCustomBreakdowns() returns custom breakdowns");
+});

devtools/client/memory/test/unit/xpcshell.ini

@@ -6,6 +6,10 @@ firefox-appdir = browser
 skip-if = toolkit == 'android' || toolkit == 'gonk'
 
 [test_action-select-snapshot.js]
+[test_action-set-breakdown.js]
+[test_action-set-breakdown-and-refresh-01.js]
+[test_action-set-breakdown-and-refresh-02.js]
 [test_action-take-census.js]
 [test_action-take-snapshot.js]
 [test_action-take-snapshot-and-census.js]
+[test_utils.js]

devtools/client/memory/utils.js

@@ -1,5 +1,7 @@
+const { Preferences } = require("resource://gre/modules/Preferences.jsm");
+const CUSTOM_BREAKDOWN_PREF = "devtools.memory.custom-breakdowns";
 const DevToolsUtils = require("devtools/shared/DevToolsUtils");
-const { snapshotState: states } = require("./constants");
+const { snapshotState: states, breakdowns } = require("./constants");
 const SAVING_SNAPSHOT_TEXT = "Saving snapshot...";
 const READING_SNAPSHOT_TEXT = "Reading snapshot...";
 const SAVING_CENSUS_TEXT = "Taking heap census...";
@@ -14,6 +16,78 @@ exports.assert = function (condition, message) {
   }
 };
 
+/**
+ * Returns an array of objects with the unique key `name`
+ * and `displayName` for each breakdown.
+ *
+ * @return {Object{name, displayName}}
+ */
+exports.getBreakdownDisplayData = function () {
+  return exports.getBreakdownNames().map(name => {
+    // If it's a preset use the display name value
+    let preset = breakdowns[name];
+    let displayName = name;
+    if (preset && preset.displayName) {
+      displayName = preset.displayName;
+    }
+    return { name, displayName };
+  });
+};
+
+/**
+ * Returns an array of the unique names for each breakdown in
+ * presets and custom pref.
+ *
+ * @return {Array<Breakdown>}
+ */
+exports.getBreakdownNames = function () {
+  let custom = exports.getCustomBreakdowns();
+  return Object.keys(Object.assign({}, breakdowns, custom));
+};
+
+/**
+ * Returns custom breakdowns defined in `devtools.memory.custom-breakdowns` pref.
+ *
+ * @return {Object}
+ */
+exports.getCustomBreakdowns = function () {
+  let customBreakdowns = Object.create(null);
+  try {
+    customBreakdowns = JSON.parse(Preferences.get(CUSTOM_BREAKDOWN_PREF)) || Object.create(null);
+  } catch (e) {
+    DevToolsUtils.reportException(
+      `String stored in "${CUSTOM_BREAKDOWN_PREF}" pref cannot be parsed by \`JSON.parse()\`.`);
+  }
+  return customBreakdowns;
+}
+
+/**
+ * Converts a breakdown preset name, like "allocationStack", and returns the
+ * spec for the breakdown. Also checks properties of keys in the `devtools.memory.custom-breakdowns`
+ * pref. If not found, returns an empty object.
+ *
+ * @param {String} name
+ * @return {Object}
+ */
+
+exports.breakdownNameToSpec = function (name) {
+  let customBreakdowns = exports.getCustomBreakdowns();
+
+  // If breakdown is already a breakdown, use it
+  if (typeof name === "object") {
+    return name;
+  }
+  // If it's in our custom breakdowns, use it
+  else if (name in customBreakdowns) {
+    return customBreakdowns[name];
+  }
+  // If breakdown name is in our presets, use that
+  else if (name in breakdowns) {
+    return breakdowns[name].breakdown;
+  }
+  return Object.create(null);
+};
+
 /**
  * Returns a string representing a readable form of the snapshot's state.
  *
@@ -21,16 +95,26 @@ exports.assert = function (condition, message) {
  * @return {String}
  */
 exports.getSnapshotStatusText = function (snapshot) {
-  switch (snapshot && snapshot.state) {
+  exports.assert((snapshot || {}).state,
+    `Snapshot must have expected state, found ${(snapshot || {}).state}.`);
+
+  switch (snapshot.state) {
     case states.SAVING:
       return SAVING_SNAPSHOT_TEXT;
     case states.SAVED:
     case states.READING:
       return READING_SNAPSHOT_TEXT;
-    case states.READ:
     case states.SAVING_CENSUS:
       return SAVING_CENSUS_TEXT;
+    // If it's read, it shouldn't have any label, as we could've cleared the
+    // census cache by changing the breakdown, and we should lazily
+    // go to SAVING_CENSUS. If it's SAVED_CENSUS, we have no status to display.
+    case states.READ:
+    case states.SAVED_CENSUS:
+      return "";
   }
+
+  DevToolsUtils.reportException(`Snapshot in unexpected state: ${snapshot.state}`);
   return "";
 }
 
@@ -66,3 +150,43 @@ exports.createSnapshot = function createSnapshot () {
     path: null,
   };
 };
+
+/**
+ * Takes two objects and compares them deeply, returning
+ * a boolean indicating if they're equal or not. Used for breakdown
+ * comparison.
+ *
+ * @param {Any} obj1
+ * @param {Any} obj2
+ * @return {Boolean}
+ */
+exports.breakdownEquals = function (obj1, obj2) {
+  let type1 = typeof obj1;
+  let type2 = typeof obj2;
+
+  // Quick checks
+  if (type1 !== type2 || (Array.isArray(obj1) !== Array.isArray(obj2))) {
+    return false;
+  }
+
+  if (obj1 === obj2) {
+    return true;
+  }
+
+  if (Array.isArray(obj1)) {
+    if (obj1.length !== obj2.length) { return false; }
+    return obj1.every((_, i) => exports.breakdownEquals(obj[1], obj2[i]));
+  }
+  else if (type1 === "object") {
+    let k1 = Object.keys(obj1);
+    let k2 = Object.keys(obj2);
+
+    if (k1.length !== k2.length) {
+      return false;
+    }
+
+    return k1.every(k => exports.breakdownEquals(obj1[k], obj2[k]));
+  }
+
+  return false;
+};

devtools/client/preferences/devtools.js

@@ -103,6 +103,8 @@ pref("devtools.debugger.ui.variables-searchbox-visible", false);
 // Enable the Memory tools
 pref("devtools.memory.enabled", false);
 
+pref("devtools.memory.custom-breakdowns", "{}");
+
 // Enable the Performance tools
 pref("devtools.performance.enabled", true);
 

devtools/shared/heapsnapshot/census-tree-node.js

@@ -63,7 +63,7 @@ CensusTreeNodeBreakdowns.internalType = function (node, breakdown, report) {
   for (let key of Object.keys(report)) {
     node.children.push(new CensusTreeNode(breakdown.then, report[key], key));
   }
-}
+};
 
 CensusTreeNodeBreakdowns.objectClass = function (node, breakdown, report) {
   node.children = [];
@@ -71,14 +71,18 @@ CensusTreeNodeBreakdowns.objectClass = function (node, breakdown, report) {
     let bd = key === "other" ? breakdown.other : breakdown.then;
     node.children.push(new CensusTreeNode(bd, report[key], key));
   }
-}
+};
 
 CensusTreeNodeBreakdowns.coarseType = function (node, breakdown, report) {
   node.children = [];
   for (let type of Object.keys(breakdown).filter(type => COARSE_TYPES.has(type))) {
     node.children.push(new CensusTreeNode(breakdown[type], report[type], type));
   }
-}
+};
+
+CensusTreeNodeBreakdowns.allocationStack = function (node, breakdown, report) {
+  node.children = [];
+};
 
 function sortByBytes (a, b) {
   return (b.bytes || 0) - (a.bytes || 0);