Hierarchical data (#8842)

* Initial HierarchicalDataProvider for TreeGrid
* Initial in-memory hierarchical data implementation
* TreeGrid declarative support

Fixes #8611, Fixes #8620

client/src/main/java/com/vaadin/client/connectors/data/DataCommunicatorConnector.java

@@ -78,6 +78,16 @@ public class DataCommunicatorConnector extends AbstractExtensionConnector {
                                 updateRowData(data.getObject(i));
                             }
                         }
+
+                        @Override
+                        public void insertRows(int firstRowIndex, int count) {
+                            insertRowData(firstRowIndex, count);
+                        }
+
+                        @Override
+                        public void removeRows(int firstRowIndex, int count) {
+                            removeRowData(firstRowIndex, count);
+                        }
                     });
         }
 
@@ -85,7 +95,8 @@ public class DataCommunicatorConnector extends AbstractExtensionConnector {
         protected void requestRows(int firstRowIndex, int numberOfRows,
                 RequestRowsCallback<JsonObject> callback) {
             getRpcProxy(DataRequestRpc.class).requestRows(firstRowIndex,
-                    numberOfRows, 0, 0);
+                    numberOfRows, getCachedRange().getStart(),
+                    getCachedRange().length());
 
             JsonArray dropped = Json.createArray();
             int i = 0;
@@ -112,13 +123,16 @@ public class DataCommunicatorConnector extends AbstractExtensionConnector {
         /**
          * Updates row data based on row key.
          *
-         * @param row
+         * @param rowData
          *            new row object
          */
-        protected void updateRowData(JsonObject row) {
-            int index = indexOfKey(getRowKey(row));
+        protected void updateRowData(JsonObject rowData) {
+            int index = indexOfKey(getRowKey(rowData));
             if (index >= 0) {
-                setRowData(index, Collections.singletonList(row));
+                JsonObject oldRowData = getRow(index);
+                onRowDataUpdate(rowData, oldRowData);
+
+                setRowData(index, Collections.singletonList(rowData));
             }
         }
     }
@@ -135,6 +149,25 @@ public class DataCommunicatorConnector extends AbstractExtensionConnector {
         }
     }
 
+    /**
+     * Called row updates from server side.
+     * <p>
+     * This method exists for making it possible to copy data from the old
+     * object to the new one, if e.g. some data is not available in the server
+     * side when doing updates and would be missed otherwise.
+     *
+     * @param newRowData
+     *            the new row data
+     * @param oldRowData
+     *            the previous row data
+     *
+     * @since 8.1
+     */
+    protected void onRowDataUpdate(JsonObject newRowData,
+            JsonObject oldRowData) {
+        // NOOP, see overrides for concrete use cases
+    }
+
     @Override
     public DataCommunicatorState getState() {
         return (DataCommunicatorState) super.getState();

client/src/main/java/com/vaadin/client/connectors/data/HierarchicalDataCommunicatorConnector.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * 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.
+ */
+package com.vaadin.client.connectors.data;
+
+import com.vaadin.data.provider.HierarchicalDataCommunicator;
+import com.vaadin.shared.ui.Connect;
+import com.vaadin.shared.ui.treegrid.TreeGridCommunicationConstants;
+
+import elemental.json.JsonObject;
+
+/**
+ * A connector for HierarchicalDataCommunicator class.
+ * 
+ * @author Vaadin Ltd
+ * @since
+ */
+@Connect(HierarchicalDataCommunicator.class)
+public class HierarchicalDataCommunicatorConnector
+        extends DataCommunicatorConnector {
+
+    @Override
+    protected void onRowDataUpdate(JsonObject newRowData,
+            JsonObject oldRowData) {
+        assert newRowData.hasKey(
+                TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION);
+        assert oldRowData.hasKey(
+                TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION);
+
+        /*
+         * Since server side can't know the index of a random item, any
+         * refreshItem(..) cannot know the depth. Thus need to copy it from
+         * previous item.
+         */
+        JsonObject hierarchyData = newRowData.getObject(
+                TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION);
+        if (!hierarchyData.hasKey(TreeGridCommunicationConstants.ROW_DEPTH)) {
+            hierarchyData.put(TreeGridCommunicationConstants.ROW_DEPTH,
+                    oldRowData
+                            .getObject(
+                                    TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION)
+                            .getNumber(
+                                    TreeGridCommunicationConstants.ROW_DEPTH));
+        }
+    }
+
+}

client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java

@@ -34,7 +34,6 @@ import com.vaadin.client.widget.grid.events.GridClickEvent;
 import com.vaadin.client.widget.treegrid.TreeGrid;
 import com.vaadin.client.widget.treegrid.events.TreeGridClickEvent;
 import com.vaadin.client.widgets.Grid;
-import com.vaadin.shared.data.DataCommunicatorConstants;
 import com.vaadin.shared.ui.Connect;
 import com.vaadin.shared.ui.treegrid.NodeCollapseRpc;
 import com.vaadin.shared.ui.treegrid.TreeGridCommunicationConstants;
@@ -81,7 +80,7 @@ public class TreeGridConnector extends GridConnector {
     void updateHierarchyColumn() {
         Scheduler.get().scheduleFinally(() -> {
             // Id of old hierarchy column
-            String oldHierarchyColumnId = this.hierarchyColumnId;
+            String oldHierarchyColumnId = hierarchyColumnId;
 
             // Id of new hierarchy column. Choose first when nothing explicitly
             // set
@@ -112,7 +111,7 @@ public class TreeGridConnector extends GridConnector {
                 // setRenderer() replaces DOM elements
                 getWidget().setFrozenColumnCount(getState().frozenColumnCount);
 
-                this.hierarchyColumnId = newHierarchyColumnId;
+                hierarchyColumnId = newHierarchyColumnId;
             } else {
                 Logger.getLogger(TreeGridConnector.class.getName()).warning(
                         "Couldn't find column: " + newHierarchyColumnId);
@@ -137,7 +136,9 @@ public class TreeGridConnector extends GridConnector {
                             @Override
                             public void onClick(
                                     ClickableRenderer.RendererClickEvent<JsonObject> event) {
-                                toggleCollapse(getRowKey(event.getRow()));
+                                toggleCollapse(getRowKey(event.getRow()),
+                                        event.getCell().getRowIndex(),
+                                        !isCollapsed(event.getRow()));
                                 event.stopPropagation();
                                 event.preventDefault();
                             }
@@ -187,8 +188,9 @@ public class TreeGridConnector extends GridConnector {
         return cell.getColumn().getRenderer() instanceof HierarchyRenderer;
     }
 
-    private void toggleCollapse(String rowKey) {
-        getRpcProxy(NodeCollapseRpc.class).toggleCollapse(rowKey);
+    private void toggleCollapse(String rowKey, int rowIndex, boolean collapse) {
+        getRpcProxy(NodeCollapseRpc.class).setNodeCollapsed(rowKey, rowIndex,
+                collapse);
     }
 
     /**
@@ -260,31 +262,26 @@ public class TreeGridConnector extends GridConnector {
 
                 // Hierarchy metadata
                 boolean collapsed, leaf;
-                if (event.getCell().getRow().hasKey(
+                JsonObject rowData = event.getCell().getRow();
+                if (rowData.hasKey(
                         TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION)) {
-                    JsonObject rowDescription = event.getCell().getRow()
-                            .getObject(
-                                    TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION);
-                    collapsed = rowDescription.getBoolean(
-                            TreeGridCommunicationConstants.ROW_COLLAPSED);
+                    collapsed = isCollapsed(rowData);
+                    JsonObject rowDescription = rowData.getObject(
+                            TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION);
                     leaf = rowDescription.getBoolean(
                             TreeGridCommunicationConstants.ROW_LEAF);
-
                     switch (domEvent.getKeyCode()) {
                     case KeyCodes.KEY_RIGHT:
-                        if (!leaf) {
-                            if (collapsed) {
-                                toggleCollapse(
-                                        event.getCell().getRow().getString(
-                                                DataCommunicatorConstants.KEY));
-                            }
+                        if (!leaf && collapsed) {
+                            toggleCollapse(getRowKey(rowData),
+                                    event.getCell().getRowIndex(), true);
                         }
                         break;
                     case KeyCodes.KEY_LEFT:
                         if (!collapsed) {
                             // collapse node
-                            toggleCollapse(event.getCell().getRow()
-                                    .getString(DataCommunicatorConstants.KEY));
+                            toggleCollapse(getRowKey(rowData),
+                                    event.getCell().getRowIndex(), false);
                         }
                         break;
                     }
@@ -294,4 +291,24 @@ public class TreeGridConnector extends GridConnector {
             }
         }
     }
+
+    private static boolean isCollapsed(JsonObject rowData) {
+        assert rowData
+                .hasKey(TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION) : "missing hierarchy data for row "
+                        + rowData.asString();
+        return rowData
+                .getObject(
+                        TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION)
+                .getBoolean(TreeGridCommunicationConstants.ROW_COLLAPSED);
+    }
+
+    private static int getDepth(JsonObject rowData) {
+        assert rowData
+                .hasKey(TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION) : "missing hierarchy data for row "
+                        + rowData.asString();
+        return (int) rowData
+                .getObject(
+                        TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION)
+                .getNumber(TreeGridCommunicationConstants.ROW_DEPTH);
+    }
 }

server/src/main/java/com/vaadin/data/HasItems.java

@@ -87,7 +87,7 @@ public interface HasItems<T> extends Component, Serializable {
      * <pre>
      * <code>
      * HasDataProvider<String> listing = new CheckBoxGroup<>();
-     * listing.setItems(Arrays.asList("a","b"));
+     * listing.setItems("a","b");
      * ...
      *
      * Collection<String> collection = ((ListDataProvider<String>)listing.getDataProvider()).getItems();
@@ -122,7 +122,7 @@ public interface HasItems<T> extends Component, Serializable {
      * <pre>
      * <code>
      * HasDataProvider<String> listing = new CheckBoxGroup<>();
-     * listing.setItems(Arrays.asList("a","b"));
+     * listing.setItems(Stream.of("a","b"));
      * ...
      *
      * Collection<String> collection = ((ListDataProvider<String>)listing.getDataProvider()).getItems();

server/src/main/java/com/vaadin/data/HierarchyData.java

@@ -0,0 +1,263 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * 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.
+ */
+package com.vaadin.data;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Stream;
+
+/**
+ * Class for representing hierarchical data.
+ * 
+ * @author Vaadin Ltd
+ * @since 8.1
+ *
+ * @param <T>
+ *            data type
+ */
+public class HierarchyData<T> implements Serializable {
+
+    private static class HierarchyWrapper<T> implements Serializable {
+        private T item;
+        private T parent;
+        private List<T> children;
+
+        public HierarchyWrapper(T item, T parent) {
+            this.item = item;
+            this.parent = parent;
+            children = new ArrayList<>();
+        }
+
+        public T getItem() {
+            return item;
+        }
+
+        public void setItem(T item) {
+            this.item = item;
+        }
+
+        public T getParent() {
+            return parent;
+        }
+
+        public void setParent(T parent) {
+            this.parent = parent;
+        }
+
+        public List<T> getChildren() {
+            return children;
+        }
+
+        public void setChildren(List<T> children) {
+            this.children = children;
+        }
+
+        public void addChild(T child) {
+            children.add(child);
+        }
+
+        public void removeChild(T child) {
+            children.remove(child);
+        }
+    }
+
+    private final Map<T, HierarchyWrapper<T>> itemToWrapperMap;
+
+    /**
+     * Creates an initially empty hierarchical data representation to which
+     * items can be added or removed.
+     */
+    public HierarchyData() {
+        itemToWrapperMap = new LinkedHashMap<>();
+        itemToWrapperMap.put(null, new HierarchyWrapper<>(null, null));
+    }
+
+    /**
+     * Adds a data item as a child of {@code parent}. Call with {@code null} as
+     * parent to add a root level item. The given parent item must already exist
+     * in this structure, and an item can only be added to this structure once.
+     * 
+     * @param parent
+     *            the parent item for which the items are added as children
+     * @param item
+     *            the item to add
+     * @return this
+     * 
+     * @throws IllegalArgumentException
+     *             if parent is not null and not already added to this structure
+     * @throws IllegalArgumentException
+     *             if the item has already been added to this structure
+     * @throws NullPointerException
+     *             if item is null
+     */
+    public HierarchyData<T> addItem(T parent, T item) {
+        Objects.requireNonNull(item, "Item cannot be null");
+        if (parent != null && !contains(parent)) {
+            throw new IllegalArgumentException(
+                    "Parent needs to be added before children. "
+                            + "To add root items, call with parent as null");
+        }
+        if (contains(item)) {
+            throw new IllegalArgumentException(
+                    "Cannot add the same item multiple times: " + item);
+        }
+        putItem(item, parent);
+        return this;
+    }
+
+    /**
+     * Adds a list of data items as children of {@code parent}. Call with
+     * {@code null} as parent to add root level items. The given parent item
+     * must already exist in this structure, and an item can only be added to
+     * this structure once.
+     * 
+     * @param parent
+     *            the parent item for which the items are added as children
+     * @param items
+     *            the list of items to add
+     * @return this
+     * 
+     * @throws IllegalArgumentException
+     *             if parent is not null and not already added to this structure
+     * @throws IllegalArgumentException
+     *             if any of the given items have already been added to this
+     *             structure
+     * @throws NullPointerException
+     *             if any of the items are null
+     */
+    public HierarchyData<T> addItems(T parent,
+            @SuppressWarnings("unchecked") T... items) {
+        Arrays.asList(items).stream().forEach(item -> addItem(parent, item));
+        return this;
+    }
+
+    /**
+     * Adds a list of data items as children of {@code parent}. Call with
+     * {@code null} as parent to add root level items. The given parent item
+     * must already exist in this structure, and an item can only be added to
+     * this structure once.
+     * 
+     * @param parent
+     *            the parent item for which the items are added as children
+     * @param items
+     *            the collection of items to add
+     * @return this
+     * 
+     * @throws IllegalArgumentException
+     *             if parent is not null and not already added to this structure
+     * @throws IllegalArgumentException
+     *             if any of the given items have already been added to this
+     *             structure
+     * @throws NullPointerException
+     *             if any of the items are null
+     */
+    public HierarchyData<T> addItems(T parent, Collection<T> items) {
+        items.stream().forEach(item -> addItem(parent, item));
+        return this;
+    }
+
+    /**
+     * Adds data items contained in a stream as children of {@code parent}. Call
+     * with {@code null} as parent to add root level items. The given parent
+     * item must already exist in this structure, and an item can only be added
+     * to this structure once.
+     * 
+     * @param parent
+     *            the parent item for which the items are added as children
+     * @param items
+     *            stream of items to add
+     * @return this
+     * 
+     * @throws IllegalArgumentException
+     *             if parent is not null and not already added to this structure
+     * @throws IllegalArgumentException
+     *             if any of the given items have already been added to this
+     *             structure
+     * @throws NullPointerException
+     *             if any of the items are null
+     */
+    public HierarchyData<T> addItems(T parent, Stream<T> items) {
+        items.forEach(item -> addItem(parent, item));
+        return this;
+    }
+
+    /**
+     * Remove a given item from this structure. Additionally, this will
+     * recursively remove any descendants of the item.
+     * 
+     * @param item
+     *            the item to remove, or null to clear all data
+     * @return this
+     * 
+     * @throws IllegalArgumentException
+     *             if the item does not exist in this structure
+     */
+    public HierarchyData<T> removeItem(T item) {
+        if (!contains(item)) {
+            throw new IllegalArgumentException(
+                    "Item '" + item + "' not in the hierarchy");
+        }
+        new ArrayList<>(getChildren(item)).forEach(child -> removeItem(child));
+        itemToWrapperMap.get(itemToWrapperMap.get(item).getParent())
+                .removeChild(item);
+        return this;
+    }
+
+    /**
+     * Clear all items from this structure. Shorthand for calling
+     * {@link #removeItem(Object)} with null.
+     */
+    public void clear() {
+        removeItem(null);
+    }
+
+    /**
+     * Get the immediate child items for the given item.
+     * 
+     * @param item
+     *            the item for which to retrieve child items for, null to
+     *            retrieve all root items
+     * @return a list of child items for the given item
+     * 
+     * @throws IllegalArgumentException
+     *             if the item does not exist in this structure
+     */
+    public List<T> getChildren(T item) {
+        if (!contains(item)) {
+            throw new IllegalArgumentException(
+                    "Item '" + item + "' not in the hierarchy");
+        }
+        return itemToWrapperMap.get(item).getChildren();
+    }
+
+    private boolean contains(T item) {
+        return itemToWrapperMap.containsKey(item);
+    }
+
+    private void putItem(T item, T parent) {
+        HierarchyWrapper<T> wrappedItem = new HierarchyWrapper<>(item, parent);
+        if (itemToWrapperMap.containsKey(parent)) {
+            itemToWrapperMap.get(parent).addChild(item);
+        }
+        itemToWrapperMap.put(item, wrappedItem);
+    }
+}

server/src/main/java/com/vaadin/data/provider/AbstractHierarchicalDataProvider.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * 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.
+ */
+package com.vaadin.data.provider;
+
+/**
+ * Abstract hierarchical data provider implementation which takes care of item
+ * refreshes and associated events.
+ *
+ * @author Vaadin Ltd
+ * @since 8.1
+ *
+ * @param <T>
+ *            data type
+ * @param <F>
+ *            filter type
+ */
+public abstract class AbstractHierarchicalDataProvider<T, F>
+        extends AbstractDataProvider<T, F>
+        implements HierarchicalDataProvider<T, F> {
+
+}

server/src/main/java/com/vaadin/data/provider/DataCommunicator.java

@@ -66,15 +66,13 @@ public class DataCommunicator<T> extends AbstractExtension {
         @Override
         public void requestRows(int firstRowIndex, int numberOfRows,
                 int firstCachedRowIndex, int cacheSize) {
-            pushRows = Range.withLength(firstRowIndex, numberOfRows);
-            markAsDirty();
+            onRequestRows(firstRowIndex, numberOfRows, firstCachedRowIndex,
+                    cacheSize);
         }
 
         @Override
         public void dropRows(JsonArray keys) {
-            for (int i = 0; i < keys.length(); ++i) {
-                handler.dropActiveData(keys.getString(i));
-            }
+            onDropRows(keys);
         }
     }
 
@@ -190,11 +188,11 @@ public class DataCommunicator<T> extends AbstractExtension {
     private final ActiveDataHandler handler = new ActiveDataHandler();
 
     /** Empty default data provider */
-    private DataProvider<T, ?> dataProvider = new CallbackDataProvider<>(
+    protected DataProvider<T, ?> dataProvider = new CallbackDataProvider<>(
             q -> Stream.empty(), q -> 0);
     private final DataKeyMapper<T> keyMapper;
 
-    private boolean reset = false;
+    protected boolean reset = false;
     private final Set<T> updatedData = new HashSet<>();
     private int minPushSize = 40;
     private Range pushRows = Range.withLength(0, minPushSize);
@@ -223,6 +221,72 @@ public class DataCommunicator<T> extends AbstractExtension {
         detachDataProviderListener();
     }
 
+    /**
+     * Set the range of rows to push for next response.
+     *
+     * @param pushRows
+     */
+    protected void setPushRows(Range pushRows) {
+        this.pushRows = pushRows;
+    }
+
+    /**
+     * Get the current range of rows to push in the next response.
+     *
+     * @return the range of rows to push
+     */
+    protected Range getPushRows() {
+        return pushRows;
+    }
+
+    /**
+     * Get the object used for filtering in this data communicator.
+     * 
+     * @return the filter object of this data communicator
+     */
+    protected Object getFilter() {
+        return filter;
+    }
+
+    /**
+     * Get the client rpc interface for this data communicator.
+     * 
+     * @return the client rpc interface for this data communicator
+     */
+    protected DataCommunicatorClientRpc getClientRpc() {
+        return rpc;
+    }
+
+    /**
+     * Request the given rows to be available on the client side.
+     *
+     * @param firstRowIndex
+     *            the index of the first requested row
+     * @param numberOfRows
+     *            the number of requested rows
+     * @param firstCachedRowIndex
+     *            the index of the first cached row
+     * @param cacheSize
+     *            the number of cached rows
+     */
+    protected void onRequestRows(int firstRowIndex, int numberOfRows,
+            int firstCachedRowIndex, int cacheSize) {
+        setPushRows(Range.withLength(firstRowIndex, numberOfRows));
+        markAsDirty();
+    }
+
+    /**
+     * Triggered when rows have been dropped from the client side cache.
+     *
+     * @param keys
+     *            the keys of the rows that have been dropped
+     */
+    protected void onDropRows(JsonArray keys) {
+        for (int i = 0; i < keys.length(); ++i) {
+            handler.dropActiveData(keys.getString(i));
+        }
+    }
+
     /**
      * Initially and in the case of a reset all data should be pushed to the
      * client.
@@ -231,6 +295,16 @@ public class DataCommunicator<T> extends AbstractExtension {
     public void beforeClientResponse(boolean initial) {
         super.beforeClientResponse(initial);
 
+        sendDataToClient(initial);
+    }
+
+    /**
+     * Send the needed data and updates to the client side.
+     *
+     * @param initial
+     *            {@code true} if initial data load, {@code false} if not
+     */
+    protected void sendDataToClient(boolean initial) {
         if (getDataProvider() == null) {
             return;
         }
@@ -241,9 +315,10 @@ public class DataCommunicator<T> extends AbstractExtension {
             rpc.reset(dataProviderSize);
         }
 
-        if (!pushRows.isEmpty()) {
-            int offset = pushRows.getStart();
-            int limit = pushRows.length();
+        Range requestedRows = getPushRows();
+        if (!requestedRows.isEmpty()) {
+            int offset = requestedRows.getStart();
+            int limit = requestedRows.length();
 
             @SuppressWarnings({ "rawtypes", "unchecked" })
             Stream<T> rowsToPush = getDataProvider().fetch(new Query(offset,
@@ -261,7 +336,7 @@ public class DataCommunicator<T> extends AbstractExtension {
             rpc.updateData(dataArray);
         }
 
-        pushRows = Range.withLength(0, 0);
+        setPushRows(Range.withLength(0, 0));
         reset = false;
         updatedData.clear();
     }
@@ -342,6 +417,15 @@ public class DataCommunicator<T> extends AbstractExtension {
         return dataObject;
     }
 
+    /**
+     * Returns the active data handler.
+     *
+     * @return the active data handler
+     */
+    protected ActiveDataHandler getActiveDataHandler() {
+        return handler;
+    }
+
     /**
      * Drops data objects identified by given keys from memory. This will invoke
      * {@link DataGenerator#destroyData} for each of those objects.
@@ -400,6 +484,15 @@ public class DataCommunicator<T> extends AbstractExtension {
         updatedData.add(data);
     }
 
+    /**
+     * Returns the currently set updated data.
+     *
+     * @return the set of data that should be updated on the next response
+     */
+    protected Set<T> getUpdatedData() {
+        return updatedData;
+    }
+
     /**
      * Sets the {@link Comparator} to use with in-memory sorting.
      *
@@ -411,6 +504,15 @@ public class DataCommunicator<T> extends AbstractExtension {
         reset();
     }
 
+    /**
+     * Returns the {@link Comparator} to use with in-memory sorting.
+     *
+     * @return comparator used to sort data
+     */
+    public Comparator<T> getInMemorySorting() {
+        return inMemorySorting;
+    }
+
     /**
      * Sets the {@link QuerySortOrder}s to use with backend sorting.
      *
@@ -423,6 +525,15 @@ public class DataCommunicator<T> extends AbstractExtension {
         reset();
     }
 
+    /**
+     * Returns the {@link QuerySortOrder} to use with backend sorting.
+     *
+     * @return list of sort order information to pass to a query
+     */
+    public List<QuerySortOrder> getBackEndSorting() {
+        return backEndSorting;
+    }
+
     /**
      * Creates a {@link DataKeyMapper} to use with this DataCommunicator.
      * <p>
@@ -492,7 +603,7 @@ public class DataCommunicator<T> extends AbstractExtension {
          * (and theoretically allows to the client doesn't request more data in
          * a happy path).
          */
-        pushRows = Range.between(0, getMinPushSize());
+        setPushRows(Range.between(0, getMinPushSize()));
         if (isAttached()) {
             attachDataProviderListener();
         }

server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java

@@ -0,0 +1,390 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * 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.
+ */
+package com.vaadin.data.provider;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiConsumer;
+import java.util.logging.Logger;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import com.vaadin.data.HierarchyData;
+import com.vaadin.data.provider.HierarchyMapper.TreeLevelQuery;
+import com.vaadin.server.SerializableConsumer;
+import com.vaadin.shared.Range;
+import com.vaadin.shared.extension.datacommunicator.HierarchicalDataCommunicatorState;
+import com.vaadin.shared.ui.treegrid.TreeGridCommunicationConstants;
+
+import elemental.json.Json;
+import elemental.json.JsonArray;
+import elemental.json.JsonObject;
+
+/**
+ * Data communicator that handles requesting hierarchical data from
+ * {@link HierarchicalDataProvider} and sending it to client side.
+ *
+ * @param <T>
+ *            the bean type
+ * @author Vaadin Ltd
+ * @since
+ */
+public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
+
+    private static final Logger LOGGER = Logger
+            .getLogger(HierarchicalDataCommunicator.class.getName());
+
+    /**
+     * The amount of root level nodes to fetch and push to the client.
+     */
+    private static final int INITIAL_FETCH_SIZE = 100;
+
+    private HierarchyMapper mapper = new HierarchyMapper();
+
+    /**
+     * The captured client side cache size.
+     */
+    private int latestCacheSize = INITIAL_FETCH_SIZE;
+
+    /**
+     * Construct a new hierarchical data communicator backed by a
+     * {@link InMemoryHierarchicalDataProvider}.
+     */
+    public HierarchicalDataCommunicator() {
+        super();
+        dataProvider = new InMemoryHierarchicalDataProvider<>(
+                new HierarchyData<>());
+    }
+
+    @Override
+    protected HierarchicalDataCommunicatorState getState() {
+        return (HierarchicalDataCommunicatorState) super.getState();
+    }
+
+    @Override
+    protected HierarchicalDataCommunicatorState getState(boolean markAsDirty) {
+        return (HierarchicalDataCommunicatorState) super.getState(markAsDirty);
+    }
+
+    @Override
+    protected void sendDataToClient(boolean initial) {
+        // on purpose do not call super
+        if (getDataProvider() == null) {
+            return;
+        }
+
+        if (initial || reset) {
+            loadInitialData();
+        } else {
+            loadRequestedRows();
+        }
+
+        if (!getUpdatedData().isEmpty()) {
+            JsonArray dataArray = Json.createArray();
+            int i = 0;
+            for (T data : getUpdatedData()) {
+                dataArray.set(i++, createDataObject(data, -1));
+            }
+            getClientRpc().updateData(dataArray);
+            getUpdatedData().clear();
+        }
+    }
+
+    private void loadInitialData() {
+        int rootSize = doSizeQuery(null);
+        mapper.reset(rootSize);
+
+        if (rootSize != 0) {
+            Range initialRange = getInitialRowsToPush(rootSize);
+            assert !initialRange
+                    .isEmpty() : "Initial range should never be empty.";
+            Stream<T> rootItems = doFetchQuery(initialRange.getStart(),
+                    initialRange.length(), null);
+
+            // for now just fetching data for the root level as everything is
+            // collapsed by default
+            List<T> items = rootItems.collect(Collectors.toList());
+            List<JsonObject> dataObjects = items.stream()
+                    .map(item -> createDataObject(item, 0))
+                    .collect(Collectors.toList());
+
+            getClientRpc().reset(rootSize);
+            sendData(0, dataObjects);
+            getActiveDataHandler().addActiveData(items.stream());
+            getActiveDataHandler().cleanUp(items.stream());
+        }
+
+        setPushRows(Range.withLength(0, 0));
+        // any updated data is ignored at this point
+        getUpdatedData().clear();
+        reset = false;
+    }
+
+    private void loadRequestedRows() {
+        final Range requestedRows = getPushRows();
+        if (!requestedRows.isEmpty()) {
+            Stream<TreeLevelQuery> levelQueries = mapper
+                    .splitRangeToLevelQueries(requestedRows.getStart(),
+                            requestedRows.getEnd() - 1);
+
+            JsonObject[] dataObjects = new JsonObject[requestedRows.length()];
+            BiConsumer<JsonObject, Integer> rowDataMapper = (object,
+                    index) -> dataObjects[index
+                            - requestedRows.getStart()] = object;
+            List<T> fetchedItems = new ArrayList<>(dataObjects.length);
+
+            levelQueries.forEach(query -> {
+                List<T> results = doFetchQuery(query.startIndex, query.size,
+                        getKeyMapper().get(query.node.getParentKey()))
+                                .collect(Collectors.toList());
+                // TODO if the size differers from expected, all goes to hell
+                fetchedItems.addAll(results);
+                List<JsonObject> rowData = results.stream()
+                        .map(item -> createDataObject(item, query.depth))
+                        .collect(Collectors.toList());
+                mapper.reorderLevelQueryResultsToFlatOrdering(rowDataMapper,
+                        query, rowData);
+            });
+            verifyNoNullItems(dataObjects, requestedRows);
+
+            sendData(requestedRows.getStart(), Arrays.asList(dataObjects));
+            getActiveDataHandler().addActiveData(fetchedItems.stream());
+            getActiveDataHandler().cleanUp(fetchedItems.stream());
+        }
+
+        setPushRows(Range.withLength(0, 0));
+    }
+
+    /*
+     * Verify that there are no null objects in the array, to fail eagerly and
+     * not just on the client side.
+     */
+    private void verifyNoNullItems(JsonObject[] dataObjects,
+            Range requestedRange) {
+        List<Integer> nullItems = new ArrayList<>(0);
+        AtomicInteger indexCounter = new AtomicInteger();
+        Stream.of(dataObjects).forEach(object -> {
+            int index = indexCounter.getAndIncrement();
+            if (object == null) {
+                nullItems.add(index);
+            }
+        });
+        if (!nullItems.isEmpty()) {
+            throw new IllegalStateException("For requested rows "
+                    + requestedRange + ", there was null items for indexes "
+                    + nullItems.stream().map(Object::toString)
+                            .collect(Collectors.joining(", ")));
+        }
+    }
+
+    private JsonObject createDataObject(T item, int depth) {
+        JsonObject dataObject = getDataObject(item);
+
+        JsonObject hierarchyData = Json.createObject();
+        if (depth != -1) {
+            hierarchyData.put(TreeGridCommunicationConstants.ROW_DEPTH, depth);
+        }
+
+        boolean isLeaf = !getDataProvider().hasChildren(item);
+        if (isLeaf) {
+            hierarchyData.put(TreeGridCommunicationConstants.ROW_LEAF, true);
+        } else {
+            String key = getKeyMapper().key(item);
+            hierarchyData.put(TreeGridCommunicationConstants.ROW_COLLAPSED,
+                    mapper.isCollapsed(key));
+            hierarchyData.put(TreeGridCommunicationConstants.ROW_LEAF, false);
+        }
+
+        // add hierarchy information to row as metadata
+        dataObject.put(TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION,
+                hierarchyData);
+
+        return dataObject;
+    }
+
+    private void sendData(int startIndex, List<JsonObject> dataObjects) {
+        JsonArray dataArray = Json.createArray();
+        int i = 0;
+        for (JsonObject dataObject : dataObjects) {
+            dataArray.set(i++, dataObject);
+        }
+
+        getClientRpc().setData(startIndex, dataArray);
+    }
+
+    /**
+     * Returns the range of rows to push on initial response.
+     *
+     * @param rootLevelSize
+     *            the amount of rows on the root level
+     * @return the range of rows to push initially
+     */
+    private Range getInitialRowsToPush(int rootLevelSize) {
+        // TODO optimize initial level to avoid unnecessary requests
+        return Range.between(0, Math.min(rootLevelSize, latestCacheSize));
+    }
+
+    @SuppressWarnings({ "rawtypes", "unchecked" })
+    private Stream<T> doFetchQuery(int start, int length, T parentItem) {
+        return getDataProvider()
+                .fetch(new HierarchicalQuery(start, length, getBackEndSorting(),
+                        getInMemorySorting(), getFilter(), parentItem));
+    }
+
+    @SuppressWarnings({ "unchecked", "rawtypes" })
+    private int doSizeQuery(T parentItem) {
+        return getDataProvider()
+                .getChildCount(new HierarchicalQuery(getFilter(), parentItem));
+    }
+
+    @Override
+    protected void onRequestRows(int firstRowIndex, int numberOfRows,
+            int firstCachedRowIndex, int cacheSize) {
+        super.onRequestRows(firstRowIndex, numberOfRows, firstCachedRowIndex,
+                cacheSize);
+    }
+
+    @Override
+    protected void onDropRows(JsonArray keys) {
+        for (int i = 0; i < keys.length(); i++) {
+            // cannot drop expanded rows since the parent item is needed always
+            // when fetching more rows
+            String itemKey = keys.getString(i);
+            if (mapper.isCollapsed(itemKey)) {
+                getActiveDataHandler().dropActiveData(itemKey);
+            }
+        }
+    }
+
+    @Override
+    public HierarchicalDataProvider<T, ?> getDataProvider() {
+        return (HierarchicalDataProvider<T, ?>) super.getDataProvider();
+    }
+
+    /**
+     * Set the current hierarchical data provider for this communicator.
+     *
+     * @param dataProvider
+     *            the data provider to set, not <code>null</code>
+     * @param initialFilter
+     *            the initial filter value to use, or <code>null</code> to not
+     *            use any initial filter value
+     *
+     * @param <F>
+     *            the filter type
+     *
+     * @return a consumer that accepts a new filter value to use
+     */
+    public <F> SerializableConsumer<F> setDataProvider(
+            HierarchicalDataProvider<T, F> dataProvider, F initialFilter) {
+        return super.setDataProvider(dataProvider, initialFilter);
+    }
+
+    /**
+     * Set the current hierarchical data provider for this communicator.
+     *
+     * @param dataProvider
+     *            the data provider to set, must extend
+     *            {@link HierarchicalDataProvider}, not <code>null</code>
+     * @param initialFilter
+     *            the initial filter value to use, or <code>null</code> to not
+     *            use any initial filter value
+     *
+     * @param <F>
+     *            the filter type
+     *
+     * @return a consumer that accepts a new filter value to use
+     */
+    @Override
+    public <F> SerializableConsumer<F> setDataProvider(
+            DataProvider<T, F> dataProvider, F initialFilter) {
+        if (dataProvider instanceof HierarchicalDataProvider) {
+            return super.setDataProvider(dataProvider, initialFilter);
+        }
+        throw new IllegalArgumentException(
+                "Only " + HierarchicalDataProvider.class.getName()
+                        + " and subtypes supported.");
+    }
+
+    /**
+     * Collapses given row, removing all its subtrees.
+     *
+     * @param collapsedRowKey
+     *            the key of the row, not {@code null}
+     * @param collapsedRowIndex
+     *            the index of row to collapse
+     */
+    public void doCollapse(String collapsedRowKey, int collapsedRowIndex) {
+        if (collapsedRowIndex < 0 | collapsedRowIndex >= mapper.getTreeSize()) {
+            throw new IllegalArgumentException("Invalid row index "
+                    + collapsedRowIndex + " when tree grid size of "
+                    + mapper.getTreeSize());
+        }
+        Objects.requireNonNull(collapsedRowKey, "Row key cannot be null");
+        T collapsedItem = getKeyMapper().get(collapsedRowKey);
+        Objects.requireNonNull(collapsedItem,
+                "Cannot find item for given key " + collapsedItem);
+
+        int collapsedSubTreeSize = mapper.collapse(collapsedRowKey,
+                collapsedRowIndex);
+
+        getClientRpc().removeRows(collapsedRowIndex + 1,
+                collapsedSubTreeSize);
+        // FIXME seems like a slight overkill to do this just for refreshing
+        // expanded status
+        refresh(collapsedItem);
+    }
+
+    /**
+     * Expands the given row.
+     *
+     * @param expandedRowKey
+     *            the key of the row, not {@code null}
+     * @param expandedRowIndex
+     *            the index of the row to expand
+     */
+    public void doExpand(String expandedRowKey, final int expandedRowIndex) {
+        if (expandedRowIndex < 0 | expandedRowIndex >= mapper.getTreeSize()) {
+            throw new IllegalArgumentException("Invalid row index "
+                    + expandedRowIndex + " when tree grid size of "
+                    + mapper.getTreeSize());
+        }
+        Objects.requireNonNull(expandedRowKey, "Row key cannot be null");
+        final T expandedItem = getKeyMapper().get(expandedRowKey);
+        Objects.requireNonNull(expandedItem,
+                "Cannot find item for given key " + expandedRowKey);
+
+        final int expandedNodeSize = doSizeQuery(expandedItem);
+        if (expandedNodeSize == 0) {
+            // TODO handle 0 size -> not expandable
+            throw new IllegalStateException("Row with index " + expandedRowIndex
+                    + " returned no child nodes.");
+        }
+
+        mapper.expand(expandedRowKey, expandedRowIndex, expandedNodeSize);
+
+        // TODO optimize by sending "enough" of the expanded items directly
+        getClientRpc().insertRows(expandedRowIndex + 1, expandedNodeSize);
+        // expanded node needs to be updated to be marked as expanded
+        // FIXME seems like a slight overkill to do this just for refreshing
+        // expanded status
+        refresh(expandedItem);
+    }
+
+}

server/src/main/java/com/vaadin/data/provider/HierarchicalDataProvider.java

@@ -15,25 +15,93 @@
  */
 package com.vaadin.data.provider;
 
+import java.util.stream.Stream;
+
 /**
- * 
+ * A common interface for fetching hierarchical data from a data source, such as
+ * an in-memory collection or a backend database.
+ *
  * @author Vaadin Ltd
  * @since 8.1
- * 
+ *
  * @param <T>
+ *            data type
  * @param <F>
+ *            filter type
  */
 public interface HierarchicalDataProvider<T, F> extends DataProvider<T, F> {
 
-    public int getDepth(T item);
+    /**
+     * Get the number of immediate child data items for the parent item returned
+     * by a given query.
+     * 
+     * @param query
+     *            given query to request the count for
+     * @return the count of child data items for the data item
+     *         {@link HierarchicalQuery#getParent()}
+     * 
+     * @throws IllegalArgumentException
+     *             if the query is not of type HierarchicalQuery
+     */
+    @Override
+    public default int size(Query<T, F> query) {
+        if (query instanceof HierarchicalQuery<?, ?>) {
+            return getChildCount((HierarchicalQuery<T, F>) query);
+        }
+        throw new IllegalArgumentException(
+                "Hierarchical data provider doesn't support non-hierarchical queries");
+    }
 
-    public boolean isRoot(T item);
+    /**
+     * Fetches data from this HierarchicalDataProvider using given
+     * {@code query}. Only the immediate children of
+     * {@link HierarchicalQuery#getParent()} will be returned.
+     * 
+     * @param query
+     *            given query to request data with
+     * @return a stream of data objects resulting from the query
+     * 
+     * @throws IllegalArgumentException
+     *             if the query is not of type HierarchicalQuery
+     */
+    @Override
+    public default Stream<T> fetch(Query<T, F> query) {
+        if (query instanceof HierarchicalQuery<?, ?>) {
+            return fetchChildren((HierarchicalQuery<T, F>) query);
+        }
+        throw new IllegalArgumentException(
+                "Hierarchical data provider doesn't support non-hierarchical queries");
+    }
 
-    public T getParent(T item);
+    /**
+     * Get the number of immediate child data items for the parent item returned
+     * by a given query.
+     * 
+     * @param query
+     *            given query to request the count for
+     * @return the count of child data items for the data item
+     *         {@link HierarchicalQuery#getParent()}
+     */
+    public int getChildCount(HierarchicalQuery<T, F> query);
 
-    public boolean isCollapsed(T item);
+    /**
+     * Fetches data from this HierarchicalDataProvider using given
+     * {@code query}. Only the immediate children of
+     * {@link HierarchicalQuery#getParent()} will be returned.
+     * 
+     * @param query
+     *            given query to request data with
+     * @return a stream of data objects resulting from the query
+     */
+    public Stream<T> fetchChildren(HierarchicalQuery<T, F> query);
 
+    /**
+     * Check whether a given item has any children associated with it.
+     * 
+     * @param item
+     *            the item to check for children
+     * @return whether the given item has children
+     */
     public boolean hasChildren(T item);
 
-    public void setCollapsed(T item, boolean b);
 }

server/src/main/java/com/vaadin/data/provider/HierarchicalQuery.java

@@ -0,0 +1,93 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * 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.
+ */
+package com.vaadin.data.provider;
+
+import java.util.Comparator;
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * Immutable hierarchical query object used to request data from a backend.
+ * Contains the parent node, index limits, sorting and filtering information.
+ *
+ * @param <T>
+ *            bean type
+ * @param <F>
+ *            filter type
+ *
+ * @since 8.0
+ */
+public class HierarchicalQuery<T, F> extends Query<T, F> {
+
+    private final T parent;
+
+    /**
+     * Constructs a new hierarchical query object with given filter and parent
+     * node.
+     *
+     * @param filter
+     *            filtering for fetching; can be <code>null</code>
+     * @param parent
+     *            the hierarchical parent object, can be <code>null</code>
+     */
+    public HierarchicalQuery(F filter, T parent) {
+        super(filter);
+        this.parent = parent;
+    }
+
+    /**
+     * Constructs a new hierarchical query object with given offset, limit,
+     * sorting and filtering.
+     *
+     * @param offset
+     *            first index to fetch
+     * @param limit
+     *            fetched item count
+     * @param sortOrders
+     *            sorting order for fetching; used for sorting backends
+     * @param inMemorySorting
+     *            comparator for sorting in-memory data
+     * @param filter
+     *            filtering for fetching; can be <code>null</code>
+     * @param parent
+     *            the hierarchical parent object, can be <code>null</code>
+     */
+    public HierarchicalQuery(int offset, int limit,
+            List<QuerySortOrder> sortOrders, Comparator<T> inMemorySorting,
+            F filter, T parent) {
+        super(offset, limit, sortOrders, inMemorySorting, filter);
+        this.parent = parent;
+    }
+
+    /**
+     * Get the hierarchical parent object, can be <code>null</code>.
+     * 
+     * @return the hierarchical parent object, can be <code>null</code>
+     */
+    public T getParent() {
+        return parent;
+    }
+
+    /**
+     * Get an Optional of the hierarchical parent object.
+     * 
+     * @see #getParent()
+     * @return the result of {@link #getParent()} wrapped by an Optional
+     */
+    public Optional<T> getParentOptional() {
+        return Optional.ofNullable(parent);
+    }
+}

server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java

@@ -0,0 +1,445 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * 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.
+ */
+package com.vaadin.data.provider;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.TreeSet;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiConsumer;
+import java.util.logging.Logger;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+/**
+ * Mapper for hierarchical data.
+ * <p>
+ * Keeps track of the expanded nodes, and size of of the subtrees for each
+ * expanded node.
+ * <p>
+ * This class is framework internal implementation details, and can be changed /
+ * moved at any point. This means that you should not directly use this for
+ * anything.
+ *
+ * @author Vaadin Ltd
+ * @since
+ */
+class HierarchyMapper implements Serializable {
+
+    private static final Logger LOGGER = Logger
+            .getLogger(HierarchyMapper.class.getName());
+
+    /**
+     * A POJO that represents a query data for a certain tree level.
+     */
+    static class TreeLevelQuery { // not serializable since not stored
+        /**
+         * The tree node that the query is for. Only used for fetching parent
+         * key.
+         */
+        final TreeNode node;
+        /** The start index of the query, from 0 to level's size - 1. */
+        final int startIndex;
+        /** The number of rows to fetch. s */
+        final int size;
+        /** The depth of this node. */
+        final int depth;
+        /** The first row index in grid, including all the nodes. */
+        final int firstRowIndex;
+        /** The direct subtrees for the node that effect the indexing. */
+        final List<TreeNode> subTrees;
+
+        TreeLevelQuery(TreeNode node, int startIndex, int size, int depth,
+                int firstRowIndex, List<TreeNode> subTrees) {
+            this.node = node;
+            this.startIndex = startIndex;
+            this.size = size;
+            this.depth = depth;
+            this.firstRowIndex = firstRowIndex;
+            this.subTrees = subTrees;
+        }
+    }
+
+    /**
+     * A level in the tree, either the root level or an expanded subtree level.
+     * <p>
+     * Comparable based on the {@link #startIndex}, which is flat from 0 to data
+     * size - 1.
+     */
+    static class TreeNode implements Serializable, Comparable<TreeNode> {
+
+        /** The key for the expanded item that this is a subtree of. */
+        private final String parentKey;
+        /** The first index on this level. */
+        private int startIndex;
+        /** The last index on this level, INCLUDING subtrees. */
+        private int endIndex;
+
+        TreeNode(String parentKey, int startIndex, int size) {
+            this.parentKey = parentKey;
+            this.startIndex = startIndex;
+            endIndex = startIndex + size - 1;
+        }
+
+        TreeNode(int startIndex) {
+            parentKey = "INVALID";
+            this.startIndex = startIndex;
+        }
+
+        int getStartIndex() {
+            return startIndex;
+        }
+
+        int getEndIndex() {
+            return endIndex;
+        }
+
+        String getParentKey() {
+            return parentKey;
+        }
+
+        private void push(int offset) {
+            startIndex += offset;
+            endIndex += offset;
+        }
+
+        private void pushEnd(int offset) {
+            endIndex += offset;
+        }
+
+        @Override
+        public int compareTo(TreeNode other) {
+            return Integer.valueOf(startIndex).compareTo(other.startIndex);
+        }
+
+        @Override
+        public String toString() {
+            return "TreeNode [parent=" + parentKey + ", start=" + startIndex
+                    + ", end=" + getEndIndex() + "]";
+        }
+
+    }
+
+    /** The expanded nodes in the tree. */
+    private final TreeSet<TreeNode> nodes = new TreeSet<>();
+
+    /**
+     * Resets the tree, sets given the root level size.
+     *
+     * @param rootLevelSize
+     *            the number of items in the root level
+     */
+    public void reset(int rootLevelSize) {
+        nodes.clear();
+        nodes.add(new TreeNode(null, 0, rootLevelSize));
+    }
+
+    /**
+     * Returns the complete size of the tree, including all expanded subtrees.
+     *
+     * @return the size of the tree
+     */
+    public int getTreeSize() {
+        TreeNode rootNode = getNodeForKey(null)
+                .orElse(new TreeNode(null, 0, 0));
+        return rootNode.endIndex + 1;
+    }
+
+    /**
+     * Returns whether the node with the given is collapsed or not.
+     *
+     * @param itemKey
+     *            the key of node to check
+     * @return {@code true} if collapsed, {@code false} if expanded
+     */
+    public boolean isCollapsed(String itemKey) {
+        return !getNodeForKey(itemKey).isPresent();
+    }
+
+    /**
+     * Return the depth of expanded node's subtree.
+     * <p>
+     * The root node depth is 0.
+     *
+     * @param expandedNodeKey
+     *            the item key of the expanded node
+     * @return the depth of the expanded node
+     * @throws IllegalArgumentException
+     *             if the node was not expanded
+     */
+    protected int getDepth(String expandedNodeKey) {
+        Optional<TreeNode> node = getNodeForKey(expandedNodeKey);
+        if (!node.isPresent()) {
+            throw new IllegalArgumentException("No node with given key "
+                    + expandedNodeKey + " was expanded.");
+        }
+        TreeNode treeNode = node.get();
+        AtomicInteger start = new AtomicInteger(treeNode.startIndex);
+        AtomicInteger end = new AtomicInteger(treeNode.getEndIndex());
+        AtomicInteger depth = new AtomicInteger();
+        nodes.headSet(treeNode, false).descendingSet().forEach(higherNode -> {
+            if (higherNode.startIndex < start.get()
+                    && higherNode.getEndIndex() >= end.get()) {
+                start.set(higherNode.startIndex);
+                depth.incrementAndGet();
+            }
+        });
+
+        return depth.get();
+    }
+
+    /**
+     * Returns the tree node for the given expanded item key, or an empty
+     * optional if the item was not expanded.
+     *
+     * @param expandedNodeKey
+     *            the key of the item
+     * @return the tree node for the expanded item, or an empty optional if not
+     *         expanded
+     */
+    protected Optional<TreeNode> getNodeForKey(String expandedNodeKey) {
+        return nodes.stream()
+                .filter(node -> Objects.equals(node.parentKey, expandedNodeKey))
+                .findAny();
+    }
+
+    /**
+     * Expands the node in the given index and with the given key.
+     *
+     * @param expanedRowKey
+     *            the key of the expanded item
+     * @param expandedRowIndex
+     *            the index of the expanded item
+     * @param expandedNodeSize
+     *            the size of the subtree of the expanded node
+     * @throws IllegalStateException
+     *             if the node was expanded already
+     */
+    protected void expand(String expanedRowKey, int expandedRowIndex,
+            int expandedNodeSize) {
+        if (expandedNodeSize < 1) {
+            throw new IllegalArgumentException(
+                    "The expanded node's size cannot be less than 1, was "
+                            + expandedNodeSize);
+        }
+        TreeNode newNode = new TreeNode(expanedRowKey, expandedRowIndex + 1,
+                expandedNodeSize);
+
+        boolean added = nodes.add(newNode);
+        if (!added) {
+            throw new IllegalStateException("Node in index " + expandedRowIndex
+                    + " was expanded already.");
+        }
+
+        // push end indexes for parent nodes
+        List<TreeNode> updated = nodes.headSet(newNode, false).stream()
+                .filter(node -> node.getEndIndex() >= expandedRowIndex)
+                .collect(Collectors.toList());
+        nodes.removeAll(updated);
+        updated.stream().forEach(node -> node.pushEnd(expandedNodeSize));
+        nodes.addAll(updated);
+
+        // push start and end indexes for later nodes
+        updated = nodes.tailSet(newNode, false).stream()
+                .collect(Collectors.toList());
+        nodes.removeAll(updated);
+        updated.stream().forEach(node -> node.push(expandedNodeSize));
+        nodes.addAll(updated);
+    }
+
+    /**
+     * Collapses the node in the given index.
+     *
+     * @param key
+     *            the key of the collapsed item
+     * @param collapsedRowIndex
+     *            the index of the collapsed item
+     * @return the size of the complete subtree that was collapsed
+     * @throws IllegalStateException
+     *             if the node was not collapsed, or if the given key is not the
+     *             same as it was when the node has been expanded
+     */
+    protected int collapse(String key, int collapsedRowIndex) {
+        Objects.requireNonNull(key,
+                "The key for the item to collapse cannot be null.");
+        TreeNode collapsedNode = nodes
+                .ceiling(new TreeNode(collapsedRowIndex + 1));
+        if (collapsedNode == null
+                || collapsedNode.startIndex != collapsedRowIndex + 1) {
+            throw new IllegalStateException(
+                    "Could not find expanded node for index "
+                            + collapsedRowIndex + ", node was not collapsed");
+        }
+        if (!Objects.equals(key, collapsedNode.parentKey)) {
+            throw new IllegalStateException("The expected parent key " + key
+                    + " is different for the collapsed node " + collapsedNode);
+        }
+
+        // remove complete subtree
+        AtomicInteger removedSubTreeSize = new AtomicInteger(
+                collapsedNode.getEndIndex() - collapsedNode.startIndex + 1);
+        nodes.tailSet(collapsedNode, false).removeIf(
+                node -> node.startIndex <= collapsedNode.getEndIndex());
+
+        final int offset = -1 * removedSubTreeSize.get();
+        // adjust parent end indexes
+        List<TreeNode> updated = nodes.headSet(collapsedNode, false).stream()
+                .filter(node -> node.getEndIndex() >= collapsedRowIndex)
+                .collect(Collectors.toList());
+        nodes.removeAll(updated);
+        updated.stream().forEach(node -> node.pushEnd(offset));
+        nodes.addAll(updated);
+
+        // adjust start and end indexes for latter nodes
+        updated = nodes.tailSet(collapsedNode, false).stream()
+                .collect(Collectors.toList());
+        nodes.removeAll(updated);
+        updated.stream().forEach(node -> node.push(offset));
+        nodes.addAll(updated);
+
+        nodes.remove(collapsedNode);
+
+        return removedSubTreeSize.get();
+    }
+
+    /**
+     * Splits the given range into queries per tree level.
+     *
+     * @param firstRow
+     *            the first row to fetch
+     * @param lastRow
+     *            the last row to fetch
+     * @return a stream of query data per level
+     * @see #reorderLevelQueryResultsToFlatOrdering(BiConsumer, TreeLevelQuery,
+     *      List)
+     */
+    protected Stream<TreeLevelQuery> splitRangeToLevelQueries(
+            final int firstRow, final int lastRow) {
+        return nodes.stream()
+                // filter to parts intersecting with the range
+                .filter(node -> node.startIndex <= lastRow
+                        && firstRow <= node.getEndIndex())
+                // split into queries per level with level based indexing
+                .map(node -> {
+
+                    // calculate how subtrees effect indexing and size
+                    int depth = getDepth(node.parentKey);
+                    List<TreeNode> directSubTrees = nodes.tailSet(node, false)
+                            .stream()
+                            // find subtrees
+                            .filter(subTree -> node.startIndex < subTree
+                                    .getEndIndex()
+                                    && subTree.startIndex < node.getEndIndex())
+                            // filter to direct subtrees
+                            .filter(subTree -> getDepth(
+                                    subTree.parentKey) == (depth + 1))
+                            .collect(Collectors.toList());
+                    // first intersecting index in flat order
+                    AtomicInteger firstIntersectingRowIndex = new AtomicInteger(
+                            Math.max(node.startIndex, firstRow));
+                    // last intersecting index in flat order
+                    final int lastIntersectingRowIndex = Math
+                            .min(node.getEndIndex(), lastRow);
+                    // start index for this level
+                    AtomicInteger start = new AtomicInteger(
+                            firstIntersectingRowIndex.get() - node.startIndex);
+                    // how many nodes should be fetched for this level
+                    AtomicInteger size = new AtomicInteger(
+                            lastIntersectingRowIndex
+                                    - firstIntersectingRowIndex.get() + 1);
+
+                    // reduce subtrees before requested index
+                    directSubTrees.stream().filter(subtree -> subtree
+                            .getEndIndex() < firstIntersectingRowIndex.get())
+                            .forEachOrdered(subtree -> {
+                                start.addAndGet(-1 * (subtree.getEndIndex()
+                                        - subtree.startIndex + 1));
+                            });
+                    // if requested start index is in the middle of a
+                    // subtree, start is after that
+                    List<TreeNode> intersectingSubTrees = new ArrayList<>();
+                    directSubTrees.stream()
+                            .filter(subtree -> subtree.startIndex <= firstIntersectingRowIndex
+                                    .get() && firstIntersectingRowIndex
+                                            .get() <= subtree.getEndIndex())
+                            .findFirst().ifPresent(subtree -> {
+                                int previous = firstIntersectingRowIndex
+                                        .getAndSet(subtree.getEndIndex() + 1);
+                                int delta = previous
+                                        - firstIntersectingRowIndex.get();
+                                start.addAndGet(subtree.startIndex - previous);
+                                size.addAndGet(delta);
+                                intersectingSubTrees.add(subtree);
+                            });
+                    // reduce size of subtrees after first row that intersect
+                    // with requested range
+                    directSubTrees.stream()
+                            .filter(subtree -> firstIntersectingRowIndex
+                                    .get() < subtree.startIndex
+                                    && subtree.endIndex <= lastIntersectingRowIndex)
+                            .forEachOrdered(subtree -> {
+                                // reduce subtree size that is part of the
+                                // requested range from query size
+                                size.addAndGet(
+                                        -1 * (Math.min(subtree.getEndIndex(),
+                                                lastIntersectingRowIndex)
+                                                - subtree.startIndex + 1));
+                                intersectingSubTrees.add(subtree);
+                            });
+                    return new TreeLevelQuery(node, start.get(), size.get(),
+                            depth, firstIntersectingRowIndex.get(),
+                            intersectingSubTrees);
+
+                }).filter(query -> query.size > 0);
+
+    }
+
+    /**
+     * Merges the tree level query results into flat grid ordering.
+     *
+     * @param rangePositionCallback
+     *            the callback to place the results into
+     * @param query
+     *            the query data for the results
+     * @param results
+     *            the results to reorder
+     * @param <T>
+     *            the type of the results
+     */
+    protected <T> void reorderLevelQueryResultsToFlatOrdering(
+            BiConsumer<T, Integer> rangePositionCallback, TreeLevelQuery query,
+            List<T> results) {
+        AtomicInteger nextPossibleIndex = new AtomicInteger(
+                query.firstRowIndex);
+        for (T item : results) {
+            // search for any intersecting subtrees and push index if necessary
+            query.subTrees.stream().filter(
+                    subTree -> subTree.startIndex <= nextPossibleIndex.get()
+                            && nextPossibleIndex.get() <= subTree.getEndIndex())
+                    .findAny().ifPresent(intersecting -> {
+                        nextPossibleIndex.addAndGet(intersecting.getEndIndex()
+                                - intersecting.startIndex + 1);
+                        query.subTrees.remove(intersecting);
+                    });
+            rangePositionCallback.accept(item,
+                    nextPossibleIndex.getAndIncrement());
+        }
+    }
+
+}

server/src/main/java/com/vaadin/data/provider/InMemoryHierarchicalDataProvider.java

@@ -0,0 +1,235 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * 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.
+ */
+package com.vaadin.data.provider;
+
+import java.util.Comparator;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.stream.Stream;
+
+import com.vaadin.data.HierarchyData;
+import com.vaadin.data.ValueProvider;
+import com.vaadin.server.SerializableComparator;
+import com.vaadin.server.SerializableFunction;
+import com.vaadin.server.SerializablePredicate;
+import com.vaadin.shared.data.sort.SortDirection;
+
+/**
+ * A {@link DataProvider} for in-memory hierarchical data.
+ *
+ * @see HierarchyData
+ *
+ * @author Vaadin Ltd
+ * @since 8.1
+ *
+ * @param <T>
+ *            data type
+ */
+public class InMemoryHierarchicalDataProvider<T> extends
+        AbstractHierarchicalDataProvider<T, SerializablePredicate<T>> implements
+        ConfigurableFilterDataProvider<T, SerializablePredicate<T>, SerializablePredicate<T>> {
+
+    private final HierarchyData<T> hierarchyData;
+
+    private SerializablePredicate<T> filter = null;
+
+    private SerializableComparator<T> sortOrder = null;
+
+    /**
+     * Constructs a new InMemoryHierarchicalDataProvider.
+     * <p>
+     * All changes made to the given HierarchyData object will also be visible
+     * through this data provider.
+     * 
+     * @param hierarchyData
+     *            the backing HierarchyData for this provider
+     */
+    public InMemoryHierarchicalDataProvider(HierarchyData<T> hierarchyData) {
+        this.hierarchyData = hierarchyData;
+    }
+
+    /**
+     * Return the underlying hierarchical data of this provider.
+     * 
+     * @return the underlying data of this provider
+     */
+    public HierarchyData<T> getData() {
+        return hierarchyData;
+    }
+
+    @Override
+    public boolean isInMemory() {
+        return true;
+    }
+
+    @Override
+    public boolean hasChildren(T item) {
+        return !hierarchyData.getChildren(item).isEmpty();
+    }
+
+    @Override
+    public int getChildCount(
+            HierarchicalQuery<T, SerializablePredicate<T>> query) {
+        return (int) fetchChildren(query).count();
+    }
+
+    @Override
+    public Stream<T> fetchChildren(
+            HierarchicalQuery<T, SerializablePredicate<T>> query) {
+        Stream<T> childStream = getFilteredStream(
+                hierarchyData.getChildren(query.getParent()).stream(),
+                query.getFilter());
+
+        Optional<Comparator<T>> comparing = Stream
+                .of(query.getInMemorySorting(), sortOrder)
+                .filter(c -> c != null)
+                .reduce((c1, c2) -> c1.thenComparing(c2));
+
+        if (comparing.isPresent()) {
+            childStream = childStream.sorted(comparing.get());
+        }
+
+        return childStream.skip(query.getOffset()).limit(query.getLimit());
+    }
+
+    @Override
+    public void setFilter(SerializablePredicate<T> filter) {
+        this.filter = filter;
+        refreshAll();
+    }
+
+    /**
+     * Adds a filter to be applied to all queries. The filter will be used in
+     * addition to any filter that has been set or added previously.
+     *
+     * @see #addFilter(ValueProvider, SerializablePredicate)
+     * @see #addFilterByValue(ValueProvider, Object)
+     * @see #setFilter(SerializablePredicate)
+     *
+     * @param filter
+     *            the filter to add, not <code>null</code>
+     */
+    public void addFilter(SerializablePredicate<T> filter) {
+        Objects.requireNonNull(filter, "Filter cannot be null");
+
+        if (this.filter == null) {
+            setFilter(filter);
+        } else {
+            SerializablePredicate<T> oldFilter = this.filter;
+            setFilter(item -> oldFilter.test(item) && filter.test(item));
+        }
+    }
+
+    /**
+     * Sets the comparator to use as the default sorting for this data provider.
+     * This overrides the sorting set by any other method that manipulates the
+     * default sorting of this data provider.
+     * <p>
+     * The default sorting is used if the query defines no sorting. The default
+     * sorting is also used to determine the ordering of items that are
+     * considered equal by the sorting defined in the query.
+     *
+     * @see #setSortOrder(ValueProvider, SortDirection)
+     * @see #addSortComparator(SerializableComparator)
+     *
+     * @param comparator
+     *            a comparator to use, or <code>null</code> to clear any
+     *            previously set sort order
+     */
+    public void setSortComparator(SerializableComparator<T> comparator) {
+        sortOrder = comparator;
+        refreshAll();
+    }
+
+    /**
+     * Adds a comparator to the default sorting for this data provider. If no
+     * default sorting has been defined, then the provided comparator will be
+     * used as the default sorting. If a default sorting has been defined, then
+     * the provided comparator will be used to determine the ordering of items
+     * that are considered equal by the previously defined default sorting.
+     * <p>
+     * The default sorting is used if the query defines no sorting. The default
+     * sorting is also used to determine the ordering of items that are
+     * considered equal by the sorting defined in the query.
+     *
+     * @see #setSortComparator(SerializableComparator)
+     * @see #addSortOrder(ValueProvider, SortDirection)
+     *
+     * @param comparator
+     *            a comparator to add, not <code>null</code>
+     */
+    public void addSortComparator(SerializableComparator<T> comparator) {
+        Objects.requireNonNull(comparator, "Sort order to add cannot be null");
+        SerializableComparator<T> originalComparator = sortOrder;
+        if (originalComparator == null) {
+            setSortComparator(comparator);
+        } else {
+            setSortComparator((a, b) -> {
+                int result = originalComparator.compare(a, b);
+                if (result == 0) {
+                    result = comparator.compare(a, b);
+                }
+                return result;
+            });
+        }
+    }
+
+    @Override
+    public <C> DataProvider<T, C> withConvertedFilter(
+            SerializableFunction<C, SerializablePredicate<T>> filterConverter) {
+        Objects.requireNonNull(filterConverter,
+                "Filter converter can't be null");
+        return new DataProviderWrapper<T, C, SerializablePredicate<T>>(this) {
+
+            @Override
+            protected SerializablePredicate<T> getFilter(Query<T, C> query) {
+                return query.getFilter().map(filterConverter).orElse(null);
+            }
+
+            @Override
+            public int size(Query<T, C> t) {
+                if (t instanceof HierarchicalQuery<?, ?>) {
+                    return dataProvider.size(new HierarchicalQuery<>(
+                            t.getOffset(), t.getLimit(), t.getSortOrders(),
+                            t.getInMemorySorting(), getFilter(t),
+                            ((HierarchicalQuery<T, C>) t).getParent()));
+                }
+                throw new IllegalArgumentException(
+                        "Hierarchical data provider doesn't support non-hierarchical queries");
+            }
+
+            @Override
+            public Stream<T> fetch(Query<T, C> t) {
+                if (t instanceof HierarchicalQuery<?, ?>) {
+                    return dataProvider.fetch(new HierarchicalQuery<>(
+                            t.getOffset(), t.getLimit(), t.getSortOrders(),
+                            t.getInMemorySorting(), getFilter(t),
+                            ((HierarchicalQuery<T, C>) t).getParent()));
+                }
+                throw new IllegalArgumentException(
+                        "Hierarchical data provider doesn't support non-hierarchical queries");
+            }
+        };
+    }
+
+    private Stream<T> getFilteredStream(Stream<T> stream,
+            Optional<SerializablePredicate<T>> queryFilter) {
+        if (filter != null) {
+            stream = stream.filter(filter);
+        }
+        return queryFilter.map(stream::filter).orElse(stream);
+    }
+}

server/src/main/java/com/vaadin/data/provider/ListDataProvider.java

@@ -31,8 +31,7 @@ import com.vaadin.shared.data.sort.SortDirection;
 import com.vaadin.ui.UI;
 
 /**
- * {@link DataProvider} wrapper for {@link Collection}s. This class does not
- * actually handle the {@link Query} parameters.
+ * {@link DataProvider} wrapper for {@link Collection}s.
  *
  * @param <T>
  *            data type

server/src/main/java/com/vaadin/ui/Grid.java

@@ -2084,20 +2084,7 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
      * @see #withPropertySet(PropertySet)
      */
     public Grid() {
-        this(new PropertySet<T>() {
-            @Override
-            public Stream<PropertyDefinition<T, ?>> getProperties() {
-                // No columns configured by default
-                return Stream.empty();
-            }
-
-            @Override
-            public Optional<PropertyDefinition<T, ?>> getProperty(String name) {
-                throw new IllegalStateException(
-                        "A Grid created without a bean type class literal or a custom property set"
-                                + " doesn't support finding properties by name.");
-            }
-        });
+        this(new DataCommunicator<>());
     }
 
     /**
@@ -2116,6 +2103,32 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
         this.beanType = beanType;
     }
 
+    /**
+     * Creates a new grid with the given data communicator and without support
+     * for creating columns based on property names.
+     *
+     * @param dataCommunicator
+     *            the custom data communicator to set
+     * @see #Grid()
+     * @see #Grid(PropertySet, DataCommunicator)
+     */
+    protected Grid(DataCommunicator<T> dataCommunicator) {
+        this(new PropertySet<T>() {
+            @Override
+            public Stream<PropertyDefinition<T, ?>> getProperties() {
+                // No columns configured by default
+                return Stream.empty();
+            }
+
+            @Override
+            public Optional<PropertyDefinition<T, ?>> getProperty(String name) {
+                throw new IllegalStateException(
+                        "A Grid created without a bean type class literal or a custom property set"
+                                + " doesn't support finding properties by name.");
+            }
+        }, dataCommunicator);
+    }
+
     /**
      * Creates a grid using a custom {@link PropertySet} implementation for
      * configuring the initial columns and resolving property names for
@@ -2128,6 +2141,27 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
      *            the property set implementation to use, not <code>null</code>.
      */
     protected Grid(PropertySet<T> propertySet) {
+        this(propertySet, new DataCommunicator<>());
+    }
+
+    /**
+     * Creates a grid using a custom {@link PropertySet} implementation and
+     * custom data communicator.
+     * <p>
+     * Property set is used for configuring the initial columns and resolving
+     * property names for {@link #addColumn(String)} and
+     * {@link Column#setEditorComponent(HasValue)}.
+     *
+     * @see #withPropertySet(PropertySet)
+     *
+     * @param propertySet
+     *            the property set implementation to use, not <code>null</code>.
+     * @param dataCommunicator
+     *            the data communicator to use, not<code>null</code>
+     */
+    protected Grid(PropertySet<T> propertySet,
+            DataCommunicator<T> dataCommunicator) {
+        super(dataCommunicator);
         registerRpc(new GridServerRpcImpl());
         setDefaultHeaderRow(appendHeaderRow());
         setSelectionModel(new SingleSelectionModelImpl<>());
@@ -3844,7 +3878,7 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
         }
     }
 
-    private void readData(Element body,
+    protected void readData(Element body,
             List<DeclarativeValueProvider<T>> providers) {
         getSelectionModel().deselectAll();
         List<T> items = new ArrayList<>();
@@ -3883,8 +3917,7 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
 
         if (designContext.shouldWriteData(this)) {
             Element bodyElement = tableElement.appendElement("tbody");
-            getDataProvider().fetch(new Query<>()).forEach(
-                    item -> writeRow(bodyElement, item, designContext));
+            writeData(bodyElement, designContext);
         }
 
         if (getFooter().getRowCount() > 0) {
@@ -3893,6 +3926,11 @@ public class Grid<T> extends AbstractListing<T> implements HasComponents,
         }
     }
 
+    protected void writeData(Element body, DesignContext designContext) {
+        getDataProvider().fetch(new Query<>())
+                .forEach(item -> writeRow(body, item, designContext));
+    }
+
     private void writeRow(Element container, T item, DesignContext context) {
         Element tableRow = container.appendElement("tr");
         tableRow.attr("item", serializeDeclarativeRepresentation(item));

server/src/main/java/com/vaadin/ui/TreeGrid.java

@@ -15,21 +15,32 @@
  */
 package com.vaadin.ui;
 
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.List;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.stream.Stream;
 
+import org.jsoup.nodes.Attributes;
+import org.jsoup.nodes.Element;
+import org.jsoup.select.Elements;
+
+import com.vaadin.data.HierarchyData;
 import com.vaadin.data.ValueProvider;
 import com.vaadin.data.provider.DataProvider;
+import com.vaadin.data.provider.HierarchicalDataCommunicator;
 import com.vaadin.data.provider.HierarchicalDataProvider;
+import com.vaadin.data.provider.HierarchicalQuery;
+import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
 import com.vaadin.shared.ui.treegrid.NodeCollapseRpc;
-import com.vaadin.shared.ui.treegrid.TreeGridCommunicationConstants;
 import com.vaadin.shared.ui.treegrid.TreeGridState;
+import com.vaadin.ui.declarative.DesignAttributeHandler;
+import com.vaadin.ui.declarative.DesignContext;
+import com.vaadin.ui.declarative.DesignFormatter;
 import com.vaadin.ui.renderers.AbstractRenderer;
 import com.vaadin.ui.renderers.Renderer;
 
-import elemental.json.Json;
-import elemental.json.JsonObject;
 
 /**
  * A grid component for displaying hierarchical tabular data.
@@ -43,55 +54,130 @@ import elemental.json.JsonObject;
 public class TreeGrid<T> extends Grid<T> {
 
     public TreeGrid() {
-        super();
-
-        // Attaches hierarchy data to the row
-        addDataGenerator((item, rowData) -> {
-
-            JsonObject hierarchyData = Json.createObject();
-            hierarchyData.put(TreeGridCommunicationConstants.ROW_DEPTH,
-                    getDataProvider().getDepth(item));
-
-            boolean isLeaf = !getDataProvider().hasChildren(item);
-            if (isLeaf) {
-                hierarchyData.put(TreeGridCommunicationConstants.ROW_LEAF,
-                        true);
-            } else {
-                hierarchyData.put(TreeGridCommunicationConstants.ROW_COLLAPSED,
-                        getDataProvider().isCollapsed(item));
-                hierarchyData.put(TreeGridCommunicationConstants.ROW_LEAF,
-                        false);
-            }
-
-            // add hierarchy information to row as metadata
-            rowData.put(
-                    TreeGridCommunicationConstants.ROW_HIERARCHY_DESCRIPTION,
-                    hierarchyData);
-        });
+        super(new HierarchicalDataCommunicator<>());
 
         registerRpc(new NodeCollapseRpc() {
             @Override
-            public void toggleCollapse(String rowKey) {
-                T item = getDataCommunicator().getKeyMapper().get(rowKey);
-                TreeGrid.this.toggleCollapse(item);
+            public void setNodeCollapsed(String rowKey, int rowIndex,
+                    boolean collapse) {
+                if (collapse) {
+                    getDataCommunicator().doCollapse(rowKey, rowIndex);
+                } else {
+                    getDataCommunicator().doExpand(rowKey, rowIndex);
+                }
             }
         });
     }
 
-    // TODO: construct a "flat" in memory hierarchical data provider?
+    /**
+     * Sets the data items of this component provided as a collection.
+     * <p>
+     * The provided items are wrapped into a
+     * {@link InMemoryHierarchicalDataProvider} backed by a flat
+     * {@link HierarchyData} structure. The data provider instance is used as a
+     * parameter for the {@link #setDataProvider(DataProvider)} method. It means
+     * that the items collection can be accessed later on via
+     * {@link InMemoryHierarchicalDataProvider#getData()}:
+     *
+     * <pre>
+     * <code>
+     * TreeGrid<String> treeGrid = new TreeGrid<>();
+     * treeGrid.setItems(Arrays.asList("a","b"));
+     * ...
+     *
+     * HierarchyData<String> data = ((InMemoryHierarchicalDataProvider<String>)treeGrid.getDataProvider()).getData();
+     * </code>
+     * </pre>
+     * <p>
+     * The returned HierarchyData instance may be used as-is to add, remove or
+     * modify items in the hierarchy. These modifications to the object are not
+     * automatically reflected back to the TreeGrid. Items modified should be
+     * refreshed with {@link HierarchicalDataProvider#refreshItem(Object)} and
+     * when adding or removing items
+     * {@link HierarchicalDataProvider#refreshAll()} should be called.
+     *
+     * @param items
+     *            the data items to display, not null
+     */
     @Override
     public void setItems(Collection<T> items) {
-        throw new UnsupportedOperationException("Not implemented");
+        Objects.requireNonNull(items, "Given collection may not be null");
+        setDataProvider(new InMemoryHierarchicalDataProvider<>(
+                new HierarchyData<T>().addItems(null, items)));
     }
 
+    /**
+     * Sets the data items of this component provided as a stream.
+     * <p>
+     * The provided items are wrapped into a
+     * {@link InMemoryHierarchicalDataProvider} backed by a flat
+     * {@link HierarchyData} structure. The data provider instance is used as a
+     * parameter for the {@link #setDataProvider(DataProvider)} method. It means
+     * that the items collection can be accessed later on via
+     * {@link InMemoryHierarchicalDataProvider#getData()}:
+     *
+     * <pre>
+     * <code>
+     * TreeGrid<String> treeGrid = new TreeGrid<>();
+     * treeGrid.setItems(Stream.of("a","b"));
+     * ...
+     *
+     * HierarchyData<String> data = ((InMemoryHierarchicalDataProvider<String>)treeGrid.getDataProvider()).getData();
+     * </code>
+     * </pre>
+     * <p>
+     * The returned HierarchyData instance may be used as-is to add, remove or
+     * modify items in the hierarchy. These modifications to the object are not
+     * automatically reflected back to the TreeGrid. Items modified should be
+     * refreshed with {@link HierarchicalDataProvider#refreshItem(Object)} and
+     * when adding or removing items
+     * {@link HierarchicalDataProvider#refreshAll()} should be called.
+     *
+     * @param items
+     *            the data items to display, not null
+     */
     @Override
     public void setItems(Stream<T> items) {
-        throw new UnsupportedOperationException("Not implemented");
+        Objects.requireNonNull(items, "Given stream may not be null");
+        setDataProvider(new InMemoryHierarchicalDataProvider<>(
+                new HierarchyData<T>().addItems(null, items)));
     }
 
+    /**
+     * Sets the data items of this listing.
+     * <p>
+     * The provided items are wrapped into a
+     * {@link InMemoryHierarchicalDataProvider} backed by a flat
+     * {@link HierarchyData} structure. The data provider instance is used as a
+     * parameter for the {@link #setDataProvider(DataProvider)} method. It means
+     * that the items collection can be accessed later on via
+     * {@link InMemoryHierarchicalDataProvider#getData()}:
+     *
+     * <pre>
+     * <code>
+     * TreeGrid<String> treeGrid = new TreeGrid<>();
+     * treeGrid.setItems("a","b");
+     * ...
+     *
+     * HierarchyData<String> data = ((InMemoryHierarchicalDataProvider<String>)treeGrid.getDataProvider()).getData();
+     * </code>
+     * </pre>
+     * <p>
+     * The returned HierarchyData instance may be used as-is to add, remove or
+     * modify items in the hierarchy. These modifications to the object are not
+     * automatically reflected back to the TreeGrid. Items modified should be
+     * refreshed with {@link HierarchicalDataProvider#refreshItem(Object)} and
+     * when adding or removing items
+     * {@link HierarchicalDataProvider#refreshAll()} should be called.
+     *
+     * @param items
+     *            the data items to display, not null
+     */
     @Override
-    public void setItems(T... items) {
-        throw new UnsupportedOperationException("Not implemented");
+    public void setItems(@SuppressWarnings("unchecked") T... items) {
+        Objects.requireNonNull(items, "Given items may not be null");
+        setDataProvider(new InMemoryHierarchicalDataProvider<>(
+                new HierarchyData<T>().addItems(null, items)));
     }
 
     @Override
@@ -138,34 +224,97 @@ public class TreeGrid<T> extends Grid<T> {
         return (TreeGridState) super.getState(markAsDirty);
     }
 
-    /**
-     * Toggle the expansion of an item in this grid. If the item is already
-     * expanded, it will be collapsed.
-     * <p>
-     * Toggling expansion on a leaf item in the hierarchy will have no effect.
-     *
-     * @param item
-     *            the item to toggle expansion for
-     */
-    public void toggleCollapse(T item) {
-        getDataProvider().setCollapsed(item,
-                !getDataProvider().isCollapsed(item));
-        getDataCommunicator().reset();
+    @Override
+    public HierarchicalDataCommunicator<T> getDataCommunicator() {
+        return (HierarchicalDataCommunicator<T>) super.getDataCommunicator();
     }
 
     @Override
     public HierarchicalDataProvider<T, ?> getDataProvider() {
-        DataProvider<T, ?> dataProvider = super.getDataProvider();
-        // FIXME DataCommunicator by default has a CallbackDataProvider if no
-        // DataProvider is set, resulting in a class cast exception if we don't
-        // check it here.
+        if (!(super.getDataProvider() instanceof HierarchicalDataProvider)) {
+            return null;
+        }
+        return (HierarchicalDataProvider<T, ?>) super.getDataProvider();
+    }
 
-        // Once fixed, remove this method from the exclude list in
-        // StateGetDoesNotMarkDirtyTest
-        if (!(dataProvider instanceof HierarchicalDataProvider)) {
-            throw new IllegalStateException("No data provider has been set.");
+    @Override
+    protected void doReadDesign(Element design, DesignContext context) {
+        super.doReadDesign(design, context);
+        Attributes attrs = design.attributes();
+        if (attrs.hasKey("hierarchy-column")) {
+            setHierarchyColumn(DesignAttributeHandler
+                    .readAttribute("hierarchy-column", attrs, String.class));
+        }
+    }
+
+    @Override
+    protected void readData(Element body,
+            List<DeclarativeValueProvider<T>> providers) {
+        getSelectionModel().deselectAll();
+        List<T> selectedItems = new ArrayList<>();
+        HierarchyData<T> data = new HierarchyData<T>();
+
+        for (Element row : body.children()) {
+            T item = deserializeDeclarativeRepresentation(row.attr("item"));
+            T parent = null;
+            if (row.hasAttr("parent")) {
+                parent = deserializeDeclarativeRepresentation(
+                        row.attr("parent"));
+            }
+            data.addItem(parent, item);
+            if (row.hasAttr("selected")) {
+                selectedItems.add(item);
+            }
+            Elements cells = row.children();
+            int i = 0;
+            for (Element cell : cells) {
+                providers.get(i).addValue(item, cell.html());
+                i++;
+            }
+        }
+
+        setDataProvider(new InMemoryHierarchicalDataProvider<>(data));
+        selectedItems.forEach(getSelectionModel()::select);
+    }
+
+    @Override
+    protected void doWriteDesign(Element design, DesignContext designContext) {
+        super.doWriteDesign(design, designContext);
+        if (getColumnByInternalId(getState(false).hierarchyColumnId) != null) {
+            String hierarchyColumn = getColumnByInternalId(
+                    getState(false).hierarchyColumnId).getId();
+            DesignAttributeHandler.writeAttribute("hierarchy-column",
+                    design.attributes(), hierarchyColumn, null, String.class,
+                    designContext);
+        }
+    }
+
+    @Override
+    protected void writeData(Element body, DesignContext designContext) {
+        getDataProvider().fetch(new HierarchicalQuery<>(null, null))
+                .forEach(item -> writeRow(body, item, null, designContext));
+    }
+
+    private void writeRow(Element container, T item, T parent,
+            DesignContext context) {
+        Element tableRow = container.appendElement("tr");
+        tableRow.attr("item", serializeDeclarativeRepresentation(item));
+        if (parent != null) {
+            tableRow.attr("parent", serializeDeclarativeRepresentation(parent));
+        }
+        if (getSelectionModel().isSelected(item)) {
+            tableRow.attr("selected", "");
+        }
+        for (Column<T, ?> column : getColumns()) {
+            Object value = column.getValueProvider().apply(item);
+            tableRow.appendElement("td")
+                    .append(Optional.ofNullable(value).map(Object::toString)
+                            .map(DesignFormatter::encodeForTextNode)
+                            .orElse(""));
         }
-        return (HierarchicalDataProvider<T, ?>) dataProvider;
+        getDataProvider().fetch(new HierarchicalQuery<>(null, item))
+                .forEach(childItem -> writeRow(container, childItem, item,
+                        context));
     }
 
     @Override

server/src/test/java/com/vaadin/data/provider/HierarchyMapperTest.java

@@ -0,0 +1,147 @@
+package com.vaadin.data.provider;
+
+import java.util.Optional;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+import com.vaadin.data.provider.HierarchyMapper.TreeNode;
+
+public class HierarchyMapperTest {
+
+    private HierarchyMapper mapper;
+
+    @Before
+    public void setup() {
+        mapper = new HierarchyMapper();
+    }
+
+    @Test
+    public void testExpandCollapse_rootLevel_indexesUpdated() {
+        mapper.reset(3);
+        verifyRootLevel(0, 2);
+
+        mapper.expand("1", 1, 3);
+
+        verifyTreeTotalSize(6);
+        verifyRootLevel(0, 5);
+        verifyNodeExists("1", 2, 4);
+
+        mapper.expand("0", 0, 3);
+
+        verifyRootLevel(0, 8);
+        verifyNodeExists("0", 1, 3);
+        verifyNodeExists("1", 5, 7);
+        verifyTreeTotalSize(9);
+
+        mapper.collapse("0", 0);
+
+        verifyRootLevel(0, 5);
+        verifyNodeExists("1", 2, 4);
+        verifyTreeTotalSize(6);
+        verifyNoNodeExists("0");
+    }
+
+    @Test
+    public void testExpandCollapse_secondLevelLastNode_indexesUpdated() {
+        mapper.reset(3);
+        verifyRootLevel(0, 2);
+
+        mapper.expand("1", 1, 3);
+
+        verifyTreeTotalSize(6);
+        verifyRootLevel(0, 5);
+        verifyNodeExists("1", 2, 4);
+
+        mapper.expand("0", 0, 3);
+
+        verifyRootLevel(0, 8);
+        verifyNodeExists("0", 1, 3);
+        verifyNodeExists("1", 5, 7);
+        verifyTreeTotalSize(9);
+
+        mapper.expand("2", 3, 3);
+
+        verifyRootLevel(0, 11);
+        verifyNodeExists("0", 1, 6);
+        verifyNodeExists("1", 8, 10);
+        verifyNodeExists("2", 4, 6);
+        verifyTreeTotalSize(12);
+
+        mapper.collapse("2", 3);
+
+        verifyRootLevel(0, 8);
+        verifyNodeExists("0", 1, 3);
+        verifyNodeExists("1", 5, 7);
+        verifyNoNodeExists("2");
+        verifyTreeTotalSize(9);
+
+        mapper.collapse("0", 0);
+
+        verifyRootLevel(0, 5);
+        verifyNodeExists("1", 2, 4);
+        verifyNoNodeExists("0");
+        verifyTreeTotalSize(6);
+    }
+
+    @Test
+    public void testCollapse_multipleLevels_wholeSubtreeDropped() {
+        // expand hierarchy up to 3 level
+        mapper.reset(5);
+        verifyRootLevel(0, 4);
+
+        mapper.expand("1", 2, 2);
+
+        verifyRootLevel(0, 6);
+        verifyNodeExists("1", 3, 4);
+        verifyTreeTotalSize(7);
+
+        mapper.expand("2", 3, 2);
+
+        verifyRootLevel(0, 8);
+        verifyNodeExists("1", 3, 6);
+        verifyNodeExists("2", 4, 5);
+        verifyTreeTotalSize(9);
+
+        mapper.expand("3", 6, 2);
+        verifyRootLevel(0, 10);
+        verifyNodeExists("1", 3, 8);
+        verifyNodeExists("2", 4, 5);
+        verifyNodeExists("3", 7, 8);
+        verifyTreeTotalSize(11);
+
+        // collapse root level node
+        mapper.collapse("1", 2);
+        verifyRootLevel(0, 4);
+        verifyNoNodeExists("1", "2", "3");
+    }
+
+    private void verifyRootLevel(int start, int end) {
+        verifyNode(start, end, mapper.getNodeForKey(null).get());
+    }
+
+    private void verifyNodeExists(String key, int start, int end) {
+        Optional<TreeNode> node = mapper.getNodeForKey(key);
+        Assert.assertTrue("NO NODE FOUND FOR KEY: " + key, node.isPresent());
+        verifyNode(start, end, node.get());
+    }
+
+    private void verifyNoNodeExists(String... nodeKeys) {
+        for (String key : nodeKeys) {
+            Assert.assertFalse("No node should exist for key " + key,
+                    mapper.getNodeForKey(key).isPresent());
+        }
+    }
+
+    private void verifyNode(int start, int end, TreeNode node) {
+        Assert.assertEquals("Invalid start for node " + node, start,
+                node.getStartIndex());
+        Assert.assertEquals("Invalid end for node " + node, end,
+                node.getEndIndex());
+    }
+
+    private void verifyTreeTotalSize(int size) {
+        Assert.assertEquals("Invalid tree size", size, mapper.getTreeSize());
+    }
+}

server/src/test/java/com/vaadin/data/provider/InMemoryHierarchicalDataProviderTest.java

@@ -0,0 +1,271 @@
+package com.vaadin.data.provider;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import com.vaadin.data.HierarchyData;
+import com.vaadin.server.SerializablePredicate;
+
+public class InMemoryHierarchicalDataProviderTest extends
+        DataProviderTestBase<InMemoryHierarchicalDataProvider<StrBean>> {
+
+    private HierarchyData<StrBean> data;
+    private List<StrBean> flattenedData;
+    private List<StrBean> rootData;
+
+    @Override
+    public void setUp() {
+        List<StrBean> randomBeans = StrBean.generateRandomBeans(20);
+        flattenedData = new ArrayList<>();
+        rootData = new ArrayList<>();
+
+        data = new HierarchyData<>();
+        data.addItems(null, randomBeans.subList(0, 5));
+        data.addItems(randomBeans.get(0), randomBeans.subList(5, 10));
+        data.addItems(randomBeans.get(5), randomBeans.subList(10, 15));
+        data.addItems(null, randomBeans.subList(15, 20));
+
+        flattenedData.add(randomBeans.get(0));
+        flattenedData.add(randomBeans.get(5));
+        flattenedData.addAll(randomBeans.subList(10, 15));
+        flattenedData.addAll(randomBeans.subList(6, 10));
+        flattenedData.addAll(randomBeans.subList(1, 5));
+        flattenedData.addAll(randomBeans.subList(15, 20));
+
+        rootData.addAll(randomBeans.subList(0, 5));
+        rootData.addAll(randomBeans.subList(15, 20));
+
+        super.setUp();
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void hierarchyData_add_item_parent_not_in_hierarchy_throws() {
+        new HierarchyData<>().addItem(new StrBean("", 0, 0),
+                new StrBean("", 0, 0));
+    }
+
+    @Test(expected = NullPointerException.class)
+    public void hierarchyData_add_null_item_throws() {
+        new HierarchyData<>().addItem(null, null);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void hierarchyData_add_item_already_in_hierarchy_throws() {
+        StrBean bean = new StrBean("", 0, 0);
+        new HierarchyData<>().addItem(null, bean).addItem(null, bean);
+    }
+
+    @Test
+    public void hierarchyData_remove_root_item() {
+        data.removeItem(null);
+        Assert.assertTrue(data.getChildren(null).isEmpty());
+    }
+
+    @Test
+    public void hierarchyData_clear() {
+        data.clear();
+        Assert.assertTrue(data.getChildren(null).isEmpty());
+    }
+
+    @Test
+    public void setFilter() {
+        getDataProvider().setFilter(item -> item.getValue().equals("Xyz")
+                || item.getValue().equals("Baz"));
+
+        Assert.assertEquals(10, sizeWithUnfilteredQuery());
+
+        getDataProvider().setFilter(item -> !item.getValue().equals("Foo")
+                && !item.getValue().equals("Xyz"));
+
+        Assert.assertEquals(
+                "Previous filter should be replaced when setting a new one", 6,
+                sizeWithUnfilteredQuery());
+
+        getDataProvider().setFilter(null);
+
+        Assert.assertEquals("Setting filter to null should remove all filters",
+                20, sizeWithUnfilteredQuery());
+    }
+
+    @Test
+    public void addFilter() {
+        getDataProvider().addFilter(item -> item.getId() <= 10);
+        getDataProvider().addFilter(item -> item.getId() >= 5);
+        Assert.assertEquals(5, sizeWithUnfilteredQuery());
+    }
+
+    @Override
+    public void filteringListDataProvider_convertFilter() {
+        DataProvider<StrBean, String> strFilterDataProvider = getDataProvider()
+                .withConvertedFilter(
+                        text -> strBean -> strBean.getValue().contains(text));
+        Assert.assertEquals("Only one item should match 'Xyz'", 1,
+                strFilterDataProvider
+                        .size(new HierarchicalQuery<>("Xyz", null)));
+        Assert.assertEquals("No item should match 'Zyx'", 0,
+                strFilterDataProvider
+                        .size(new HierarchicalQuery<>("Zyx", null)));
+        Assert.assertEquals("Unexpected number of matches for 'Foo'", 3,
+                strFilterDataProvider
+                        .size(new HierarchicalQuery<>("Foo", null)));
+        Assert.assertEquals("No items should've been filtered out",
+                rootData.size(), strFilterDataProvider
+                        .size(new HierarchicalQuery<>(null, null)));
+    }
+
+    @Override
+    public void filteringListDataProvider_defaultFilterType() {
+        Assert.assertEquals("Only one item should match 'Xyz'", 1,
+                getDataProvider().size(new HierarchicalQuery<>(
+                        strBean -> strBean.getValue().contains("Xyz"), null)));
+        Assert.assertEquals("No item should match 'Zyx'", 0,
+                dataProvider.size(new HierarchicalQuery<>(
+                        strBean -> strBean.getValue().contains("Zyx"), null)));
+        Assert.assertEquals("Unexpected number of matches for 'Foo'", 3,
+                getDataProvider()
+                        .size(new HierarchicalQuery<>(fooFilter, null)));
+    }
+
+    @Override
+    public void testDefaultSortWithSpecifiedPostSort() {
+        Comparator<StrBean> comp = Comparator.comparing(StrBean::getValue)
+                .thenComparing(Comparator.comparing(StrBean::getId).reversed());
+        setSortOrder(QuerySortOrder.asc("value").thenDesc("id").build(), comp);
+
+        List<StrBean> list = getDataProvider()
+                .fetch(createQuery(QuerySortOrder.asc("randomNumber").build(),
+                        Comparator.comparing(StrBean::getRandomNumber), null,
+                        null))
+                .collect(Collectors.toList());
+
+        Assert.assertEquals("Sorted data and original data sizes don't match",
+                getDataProvider().fetch(new HierarchicalQuery<>(null, null))
+                        .count(),
+                list.size());
+
+        for (int i = 1; i < list.size(); ++i) {
+            StrBean prev = list.get(i - 1);
+            StrBean cur = list.get(i);
+            // Test specific sort
+            Assert.assertTrue(
+                    "Failure: " + prev.getRandomNumber() + " > "
+                            + cur.getRandomNumber(),
+                    prev.getRandomNumber() <= cur.getRandomNumber());
+
+            if (prev.getRandomNumber() == cur.getRandomNumber()) {
+                // Test default sort
+                Assert.assertTrue(
+                        prev.getValue().compareTo(cur.getValue()) <= 0);
+                if (prev.getValue().equals(cur.getValue())) {
+                    Assert.assertTrue(prev.getId() > cur.getId());
+                }
+            }
+        }
+    }
+
+    @Override
+    public void testDefaultSortWithFunction() {
+        setSortOrder(QuerySortOrder.asc("value").build(),
+                Comparator.comparing(StrBean::getValue));
+
+        List<StrBean> list = getDataProvider()
+                .fetch(new HierarchicalQuery<>(null, null))
+                .collect(Collectors.toList());
+
+        Assert.assertEquals("Sorted data and original data sizes don't match",
+                rootData.size(), list.size());
+
+        for (int i = 1; i < list.size(); ++i) {
+            StrBean prev = list.get(i - 1);
+            StrBean cur = list.get(i);
+
+            // Test default sort
+            Assert.assertTrue(prev.getValue().compareTo(cur.getValue()) <= 0);
+        }
+    }
+
+    @Override
+    public void testListContainsAllData() {
+        assertHierarchyCorrect();
+    }
+
+    @Override
+    public void testSortByComparatorListsDiffer() {
+        Comparator<StrBean> comp = Comparator.comparing(StrBean::getValue)
+                .thenComparing(StrBean::getRandomNumber)
+                .thenComparing(StrBean::getId);
+
+        List<StrBean> list = getDataProvider().fetch(
+                createQuery(QuerySortOrder.asc("value").thenAsc("randomNumber")
+                        .thenAsc("id").build(), comp, null, null))
+                .collect(Collectors.toList());
+
+        Assert.assertNotEquals("First value should not match", rootData.get(0),
+                list.get(0));
+
+        Assert.assertEquals("Sorted data and original data sizes don't match",
+                rootData.size(), list.size());
+
+        rootData.sort(comp);
+        for (int i = 0; i < rootData.size(); ++i) {
+            Assert.assertEquals("Sorting result differed", rootData.get(i),
+                    list.get(i));
+        }
+    }
+
+    @Override
+    protected InMemoryHierarchicalDataProvider<StrBean> createDataProvider() {
+        return new InMemoryHierarchicalDataProvider<>(data);
+    }
+
+    @Override
+    protected void setSortOrder(List<QuerySortOrder> sortOrder,
+            Comparator<StrBean> comp) {
+        getDataProvider().setSortComparator(comp::compare);
+    }
+
+    @Override
+    protected long sizeWithUnfilteredQuery() {
+        return getFlattenedDataFromProvider(new ArrayList<>(), null).size();
+    }
+
+    private void assertHierarchyCorrect() {
+        Assert.assertEquals(flattenedData,
+                getFlattenedData(new ArrayList<>(), null));
+        Assert.assertEquals(flattenedData,
+                getFlattenedDataFromProvider(new ArrayList<>(), null));
+    }
+
+    private List<StrBean> getFlattenedData(List<StrBean> flattened,
+            StrBean item) {
+        if (item != null) {
+            flattened.add(item);
+        }
+        data.getChildren(item)
+                .forEach(child -> getFlattenedData(flattened, child));
+        return flattened;
+    }
+
+    private List<StrBean> getFlattenedDataFromProvider(List<StrBean> flattened,
+            StrBean item) {
+        if (item != null) {
+            flattened.add(item);
+        }
+        getDataProvider().fetchChildren(new HierarchicalQuery<>(null, item))
+                .forEach(child -> getFlattenedDataFromProvider(flattened,
+                        child));
+        return flattened;
+    }
+
+    private HierarchicalQuery<StrBean, SerializablePredicate<StrBean>> createQuery(
+            List<QuerySortOrder> sortOrder, Comparator<StrBean> comp,
+            SerializablePredicate<StrBean> filter, StrBean parent) {
+        return new HierarchicalQuery<>(0, Integer.MAX_VALUE, sortOrder, comp,
+                filter, parent);
+    }
+}

server/src/test/java/com/vaadin/tests/server/ClassesSerializableTest.java

@@ -76,6 +76,7 @@ public class ClassesSerializableTest {
             "com\\.vaadin\\.buildhelpers.*", //
             "com\\.vaadin\\.util\\.EncodeUtil.*", //
             "com\\.vaadin\\.util\\.ReflectTools.*", //
+            "com\\.vaadin\\.data\\.provider\\.HierarchyMapper\\$TreeLevelQuery",
             "com\\.vaadin\\.data\\.util\\.ReflectTools.*", //
             "com\\.vaadin\\.data\\.util\\.JsonUtil.*", //
             "com\\.vaadin\\.data\\.util.BeanItemContainerGenerator.*",

server/src/test/java/com/vaadin/tests/server/component/treegrid/TreeGridDeclarativeTest.java

@@ -0,0 +1,116 @@
+package com.vaadin.tests.server.component.treegrid;
+
+import java.lang.reflect.InvocationTargetException;
+
+import org.junit.Assert;
+
+import com.vaadin.data.HierarchyData;
+import com.vaadin.data.provider.HierarchicalQuery;
+import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
+import com.vaadin.tests.data.bean.Person;
+import com.vaadin.tests.server.component.abstractlisting.AbstractListingDeclarativeTest;
+import com.vaadin.ui.TreeGrid;
+
+public class TreeGridDeclarativeTest
+        extends AbstractListingDeclarativeTest<TreeGrid> {
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public void dataSerialization() throws InstantiationException,
+            IllegalAccessException, InvocationTargetException {
+        TreeGrid<Person> grid = new TreeGrid<>();
+
+        Person person1 = createPerson("a", "last-name");
+        Person person2 = createPerson("aa", "last-name");
+        Person person3 = createPerson("ab", "last-name");
+        Person person4 = createPerson("b", "last-name");
+        Person person5 = createPerson("c", "last-name");
+        Person person6 = createPerson("ca", "last-name");
+        Person person7 = createPerson("caa", "last-name");
+
+        HierarchyData<Person> data = new HierarchyData<>();
+        data.addItems(null, person1, person4, person5);
+        data.addItems(person1, person2, person3);
+        data.addItem(person5, person6);
+        data.addItem(person6, person7);
+
+        grid.addColumn(Person::getFirstName).setCaption("First Name");
+        grid.addColumn(Person::getLastName).setId("id").setCaption("Id");
+
+        grid.setHierarchyColumn("id");
+        grid.setDataProvider(new InMemoryHierarchicalDataProvider<>(data));
+
+        String design = String.format(
+                "<%s hierarchy-column='id'><table><colgroup>"
+                        + "<col column-id='column0' sortable>"
+                        + "<col column-id='id' sortable></colgroup><thead>"
+                        + "<tr default><th plain-text column-ids='column0'>First Name</th>"
+                        + "<th plain-text column-ids='id'>Id</th></tr>"
+                        + "</thead><tbody>"
+                        + "<tr item='%s'><td>%s</td><td>%s</td></tr>"
+                        + "<tr item='%s' parent='%s'><td>%s</td><td>%s</td></tr>"
+                        + "<tr item='%s' parent='%s'><td>%s</td><td>%s</td></tr>"
+                        + "<tr item='%s'><td>%s</td><td>%s</td></tr>"
+                        + "<tr item='%s'><td>%s</td><td>%s</td></tr>"
+                        + "<tr item='%s' parent='%s'><td>%s</td><td>%s</td></tr>"
+                        + "<tr item='%s' parent='%s'><td>%s</td><td>%s</td></tr>"
+                        + "</tbody></table></%s>",
+                getComponentTag(), person1.toString(), person1.getFirstName(),
+                person1.getLastName(), person2.toString(), person1.toString(),
+                person2.getFirstName(), person2.getLastName(),
+                person3.toString(), person1.toString(), person3.getFirstName(),
+                person3.getLastName(), person4.toString(),
+                person4.getFirstName(), person4.getLastName(),
+                person5.toString(), person5.getFirstName(),
+                person5.getLastName(), person6.toString(), person5.toString(),
+                person6.getFirstName(), person6.getLastName(),
+                person7.toString(), person6.toString(), person7.getFirstName(),
+                person7.getLastName(), getComponentTag());
+
+        TreeGrid<String> readGrid = testRead(design, grid);
+        Assert.assertEquals(3, readGrid.getDataProvider()
+                .size(new HierarchicalQuery<>(null, null)));
+        Assert.assertEquals(2, readGrid.getDataProvider()
+                .size(new HierarchicalQuery<>(null, person1.toString())));
+        Assert.assertEquals(1, readGrid.getDataProvider()
+                .size(new HierarchicalQuery<>(null, person5.toString())));
+        Assert.assertEquals(1, readGrid.getDataProvider()
+                .size(new HierarchicalQuery<>(null, person6.toString())));
+        testWrite(design, grid, true);
+    }
+    
+
+    @Override
+    public void valueSerialization() throws InstantiationException,
+            IllegalAccessException, InvocationTargetException {
+        // Tested by GridDeclarativeTest
+    }
+
+    @Override
+    public void readOnlySelection() throws InstantiationException,
+            IllegalAccessException, InvocationTargetException {
+        // Tested by GridDeclarativeTest
+    }
+
+    @Override
+    protected String getComponentTag() {
+        return "vaadin-tree-grid";
+    }
+
+    @Override
+    protected Class<? extends TreeGrid> getComponentClass() {
+        return TreeGrid.class;
+    }
+
+    private Person createPerson(String name, String lastName) {
+        Person person = new Person() {
+            @Override
+            public String toString() {
+                return getFirstName() + " " + getLastName();
+            }
+        };
+        person.setFirstName(name);
+        person.setLastName(lastName);
+        return person;
+    }
+}

shared/src/main/java/com/vaadin/shared/data/DataCommunicatorClientRpc.java

@@ -56,5 +56,24 @@ public interface DataCommunicatorClientRpc extends ClientRpc {
      */
     void updateData(JsonArray data);
 
-    // TODO: Notify add / remove
+    /**
+     * Informs that new data has been inserted from the server.
+     *
+     * @param firstRowIndex
+     *            the destination index of the new row data
+     * @param count
+     *            the number of rows inserted
+     */
+    void insertRows(int firstRowIndex, int count);
+
+    /**
+     * Informs that the server has removed data.
+     *
+     * @param firstRowIndex
+     *            the index of the first removed row
+     * @param count
+     *            the number of removed rows, starting from
+     *            <code>firstRowIndex</code>
+     */
+    void removeRows(int firstRowIndex, int count);
 }

shared/src/main/java/com/vaadin/shared/extension/datacommunicator/HierarchicalDataCommunicatorState.java

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2000-2016 Vaadin Ltd.
+ *
+ * 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.
+ */
+package com.vaadin.shared.extension.datacommunicator;
+
+/**
+ * Shared state for HierarchicalDataCommunicator.
+ *
+ * @since
+ */
+public class HierarchicalDataCommunicatorState extends DataCommunicatorState {
+
+}

shared/src/main/java/com/vaadin/shared/ui/treegrid/NodeCollapseRpc.java

@@ -20,12 +20,22 @@ import com.vaadin.shared.communication.ServerRpc;
 /**
  * RPC to handle client originated collapse and expand actions on hierarchical
  * rows in TreeGrid.
- * 
+ *
  * @author Vaadin Ltd
  * @since 8.1
  */
 @FunctionalInterface
 public interface NodeCollapseRpc extends ServerRpc {
 
-    void toggleCollapse(String rowKey);
+    /**
+     * Sets the collapse state of a hierarchical row in TreeGrid.
+     *
+     * @param rowKey
+     *            the row's key
+     * @param rowIndex
+     *            index where the row is in grid (all rows)
+     * @param collapse
+     *            {@code true} to collapse, {@code false} to expand
+     */
+    void setNodeCollapsed(String rowKey, int rowIndex, boolean collapse);
 }

testbench-api/src/main/java/com/vaadin/testbench/elements/TreeGridElement.java

@@ -15,6 +15,11 @@
  */
 package com.vaadin.testbench.elements;
 
+import org.openqa.selenium.NoSuchElementException;
+import org.openqa.selenium.WebElement;
+
+import com.vaadin.testbench.By;
+
 /**
  * TestBench Element API for TreeGrid
  *
@@ -22,4 +27,111 @@ package com.vaadin.testbench.elements;
  */
 public class TreeGridElement extends GridElement {
 
+    /**
+     * Expands the row at the given index in the grid. This expects the first
+     * column to have the hierarchy data.
+     *
+     * @param rowIndex
+     *            0-based row index to expand
+     * @see #expandWithClick(int, int)
+     */
+    public void expandWithClick(int rowIndex) {
+        expandWithClick(rowIndex, 0);
+    }
+
+    /**
+     * Expands the row at the given index in the grid with the given
+     * hierarchical column index.
+     *
+     * @param rowIndex
+     *            0-based row index to expand
+     * @param hierarchyColumnIndex
+     *            0-based index of the hierarchy column
+     */
+    public void expandWithClick(int rowIndex, int hierarchyColumnIndex) {
+        if (isRowExpanded(rowIndex, hierarchyColumnIndex)) {
+            throw new IllegalStateException(
+                    "The element at row " + rowIndex + " was expanded already");
+        }
+        getExpandElement(rowIndex, hierarchyColumnIndex).click();
+    }
+
+    /**
+     * Collapses the row at the given index in the grid. This expects the first
+     * column to have the hierarchy data.
+     *
+     * @param rowIndex
+     *            0-based row index to collapse
+     * @see #collapseWithClick(int, int)
+     */
+    public void collapseWithClick(int rowIndex) {
+        collapseWithClick(rowIndex, 0);
+    }
+
+    /**
+     * Collapses the row at the given index in the grid with the given
+     * hierarchical column index.
+     *
+     * @param rowIndex
+     *            0-based row index to collapse
+     * @param hierarchyColumnIndex
+     *            0-based index of the hierarchy column
+     */
+    public void collapseWithClick(int rowIndex, int hierarchyColumnIndex) {
+        if (isRowCollapsed(rowIndex, hierarchyColumnIndex)) {
+            throw new IllegalStateException("The element at row " + rowIndex
+                    + " was collapsed already");
+        }
+        getExpandElement(rowIndex, hierarchyColumnIndex).click();
+    }
+
+    /**
+     * Returns whether the row at the given index is expanded or not.
+     *
+     * @param rowIndex
+     *            0-based row index
+     * @param hierarchyColumnIndex
+     *            0-based index of the hierarchy column
+     * @return {@code true} if expanded, {@code false} if collapsed
+     */
+    public boolean isRowExpanded(int rowIndex, int hierarchyColumnIndex) {
+        WebElement expandElement = getExpandElement(rowIndex,
+                hierarchyColumnIndex);
+        return expandElement.getAttribute("expanded") != null
+                && expandElement.getAttribute("collapsed") == null;
+    }
+
+    /**
+     * Returns whether the row at the given index is collapsed or not.
+     *
+     * @param rowIndex
+     *            0-based row index
+     * @param hierarchyColumnIndex
+     *            0-based index of the hierarchy column
+     * @return {@code true} if collapsed, {@code false} if expanded
+     */
+    public boolean isRowCollapsed(int rowIndex, int hierarchyColumnIndex) {
+        WebElement expandElement = getExpandElement(rowIndex,
+                hierarchyColumnIndex);
+        return expandElement.getAttribute("collapsed") != null
+                && expandElement.getAttribute("expanded") == null;
+    }
+
+    /**
+     * Gets the expand/collapse element for the given row.
+     *
+     * @param rowIndex
+     *            0-based row index
+     * @param hierarchyColumnIndex
+     *            0-based index of the hierarchy column
+     * @return the {@code span} element that is clicked for expanding/collapsing
+     *         a rows
+     * @throws NoSuchElementException
+     *             if there is no expand element for this row
+     */
+    public WebElement getExpandElement(int rowIndex, int hierarchyColumnIndex) {
+        return getCell(rowIndex, hierarchyColumnIndex)
+                .findElement(By.className("v-tree-grid-expander"));
+
+    }
 }

uitest/src/main/java/com/vaadin/tests/components/treegrid/LazyHierarchicalDataProvider.java

@@ -0,0 +1,63 @@
+package com.vaadin.tests.components.treegrid;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.stream.Stream;
+
+import com.vaadin.data.provider.AbstractHierarchicalDataProvider;
+import com.vaadin.data.provider.HierarchicalQuery;
+import com.vaadin.tests.components.treegrid.TreeGridBasicFeatures.HierarchicalTestBean;
+
+public class LazyHierarchicalDataProvider
+        extends AbstractHierarchicalDataProvider<HierarchicalTestBean, Void> {
+
+    private final int nodesPerLevel;
+    private final int depth;
+
+    public LazyHierarchicalDataProvider(int nodesPerLevel, int depth) {
+        this.nodesPerLevel = nodesPerLevel;
+        this.depth = depth;
+    }
+
+    @Override
+    public int getChildCount(
+            HierarchicalQuery<HierarchicalTestBean, Void> query) {
+
+        Optional<Integer> count = query.getParentOptional()
+                .flatMap(parent -> Optional.of(Integer.valueOf(
+                        (internalHasChildren(parent) ? nodesPerLevel : 0))));
+
+        return count.orElse(nodesPerLevel);
+    }
+
+    @Override
+    public Stream<HierarchicalTestBean> fetchChildren(
+            HierarchicalQuery<HierarchicalTestBean, Void> query) {
+        final int depth = query.getParentOptional().isPresent()
+                ? query.getParent().getDepth() + 1 : 0;
+        final Optional<String> parentKey = query.getParentOptional()
+                .flatMap(parent -> Optional.of(parent.getId()));
+
+        List<HierarchicalTestBean> list = new ArrayList<>();
+        for (int i = 0; i < query.getLimit(); i++) {
+            list.add(new HierarchicalTestBean(parentKey.orElse(null), depth,
+                    i + query.getOffset()));
+        }
+        return list.stream();
+    }
+
+    @Override
+    public boolean hasChildren(HierarchicalTestBean item) {
+        return internalHasChildren(item);
+    }
+
+    private boolean internalHasChildren(HierarchicalTestBean node) {
+        return node.getDepth() < depth;
+    }
+
+    @Override
+    public boolean isInMemory() {
+        return false;
+    }
+}

uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeatures.java

@@ -1,28 +1,24 @@
 package com.vaadin.tests.components.treegrid;
 
-import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.LinkedHashMap;
-import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
-import com.vaadin.data.provider.DataProviderListener;
-import com.vaadin.data.provider.HierarchicalDataProvider;
-import com.vaadin.data.provider.Query;
-import com.vaadin.shared.Registration;
+import com.vaadin.annotations.Theme;
+import com.vaadin.annotations.Widgetset;
+import com.vaadin.data.HierarchyData;
+import com.vaadin.data.provider.DataProvider;
+import com.vaadin.data.provider.InMemoryHierarchicalDataProvider;
 import com.vaadin.tests.components.AbstractComponentTest;
-import com.vaadin.ui.MenuBar.MenuItem;
 import com.vaadin.ui.TreeGrid;
 
+@Theme("valo")
+@Widgetset("com.vaadin.DefaultWidgetSet")
 public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
 
-    private TreeGrid<TestBean> grid;
-    private TestDataProvider dataProvider = new TestDataProvider();
+    private TreeGrid<HierarchicalTestBean> grid;
+    private InMemoryHierarchicalDataProvider<HierarchicalTestBean> inMemoryDataProvider;
+    private LazyHierarchicalDataProvider lazyDataProvider;
 
     @Override
     public TreeGrid getComponent() {
@@ -36,12 +32,17 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
 
     @Override
     protected void initializeComponents() {
+        initializeDataProviders();
         grid = new TreeGrid<>();
         grid.setSizeFull();
-        grid.addColumn(TestBean::getStringValue).setId("First column");
-        grid.addColumn(TestBean::getStringValue).setId("Second column");
-        grid.setHierarchyColumn("First column");
-        grid.setDataProvider(dataProvider);
+        grid.addColumn(HierarchicalTestBean::toString).setCaption("String")
+                .setId("string");
+        grid.addColumn(HierarchicalTestBean::getDepth).setCaption("Depth")
+                .setId("depth");
+        grid.addColumn(HierarchicalTestBean::getIndex)
+                .setCaption("Index on this depth").setId("index");
+        grid.setHierarchyColumn("string");
+        grid.setDataProvider(new LazyHierarchicalDataProvider(3, 2));
 
         grid.setId("testComponent");
         addTestComponent(grid);
@@ -51,8 +52,44 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
     protected void createActions() {
         super.createActions();
 
+        createDataProviderSelect();
         createHierarchyColumnSelect();
-        createToggleCollapseSelect();
+    }
+
+    private void initializeDataProviders() {
+        HierarchyData<HierarchicalTestBean> data = new HierarchyData<>();
+
+        List<Integer> ints = Arrays.asList(0, 1, 2);
+
+        ints.stream().forEach(index -> {
+            HierarchicalTestBean bean = new HierarchicalTestBean(null, 0,
+                    index);
+            data.addItem(null, bean);
+            ints.stream().forEach(childIndex -> {
+                HierarchicalTestBean childBean = new HierarchicalTestBean(
+                        bean.getId(), 1, childIndex);
+                data.addItem(bean, childBean);
+                ints.stream()
+                        .forEach(grandChildIndex -> data.addItem(childBean,
+                                new HierarchicalTestBean(childBean.getId(), 2,
+                                        grandChildIndex)));
+            });
+        });
+
+        inMemoryDataProvider = new InMemoryHierarchicalDataProvider<>(data);
+        lazyDataProvider = new LazyHierarchicalDataProvider(3, 2);
+    }
+
+    @SuppressWarnings("unchecked")
+    private void createDataProviderSelect() {
+        @SuppressWarnings("rawtypes")
+        LinkedHashMap<String, DataProvider> options = new LinkedHashMap<>();
+        options.put("LazyHierarchicalDataProvider", lazyDataProvider);
+        options.put("InMemoryHierarchicalDataProvider", inMemoryDataProvider);
+
+        createSelectAction("Set data provider", CATEGORY_FEATURES, options,
+                "LazyHierarchicalDataProvider",
+                (treeGrid, value, data) -> treeGrid.setDataProvider(value));
     }
 
     private void createHierarchyColumnSelect() {
@@ -65,205 +102,64 @@ public class TreeGridBasicFeatures extends AbstractComponentTest<TreeGrid> {
                 (treeGrid, value, data) -> treeGrid.setHierarchyColumn(value));
     }
 
-    private void createToggleCollapseSelect() {
-        MenuItem menu = createCategory("Toggle expand", CATEGORY_FEATURES);
-        dataProvider.getAllItems().forEach(testBean -> {
-            createClickAction(testBean.getStringValue(), "Toggle expand",
-                    (grid, bean, data) -> grid.toggleCollapse(bean), testBean);
-        });
-    }
-
-    private static class TestBean {
-
-        private String stringValue;
-
-        public TestBean(String stringValue) {
-            this.stringValue = stringValue;
-        }
-
-        public String getStringValue() {
-            return stringValue;
-        }
-
-        public void setStringValue(String stringValue) {
-            this.stringValue = stringValue;
-        }
-    }
-
-    private static class TestDataProvider
-            implements HierarchicalDataProvider<TestBean, Void> {
-
-        private static class HierarchyWrapper<T> {
-            private T item;
-            private T parent;
-            private Set<T> children;
-            private boolean collapsed;
-
-            public HierarchyWrapper(T item, T parent, boolean collapsed) {
-                this.item = item;
-                this.parent = parent;
-                this.collapsed = collapsed;
-                children = new LinkedHashSet<>();
-            }
-
-            public T getItem() {
-                return item;
-            }
-
-            public void setItem(T item) {
-                this.item = item;
-            }
-
-            public T getParent() {
-                return parent;
-            }
-
-            public void setParent(T parent) {
-                this.parent = parent;
-            }
-
-            public Set<T> getChildren() {
-                return children;
-            }
-
-            public void setChildren(Set<T> children) {
-                this.children = children;
-            }
-
-            public boolean isCollapsed() {
-                return collapsed;
-            }
-
-            public void setCollapsed(boolean collapsed) {
-                this.collapsed = collapsed;
-            }
-        }
-
-        private Map<TestBean, HierarchyWrapper<TestBean>> itemToWrapperMap;
-        private Map<HierarchyWrapper<TestBean>, TestBean> wrapperToItemMap;
-        private Map<TestBean, HierarchyWrapper<TestBean>> rootNodes;
-
-        public TestDataProvider() {
-            itemToWrapperMap = new LinkedHashMap<>();
-            wrapperToItemMap = new LinkedHashMap<>();
-            rootNodes = new LinkedHashMap<>();
+    static class HierarchicalTestBean {
 
-            List<String> strings = Arrays.asList("a", "b", "c");
+        private final String id;
+        private final int depth;
+        private final int index;
 
-            strings.stream().forEach(string -> {
-                TestBean rootBean = new TestBean(string);
-
-                HierarchyWrapper<TestBean> wrappedParent = new HierarchyWrapper<>(
-                        rootBean, null, true);
-                itemToWrapperMap.put(rootBean, wrappedParent);
-                wrapperToItemMap.put(wrappedParent, rootBean);
-
-                List<TestBean> children = strings.stream().map(string2 -> {
-                    TestBean childBean = new TestBean(string + "/" + string2);
-                    HierarchyWrapper<TestBean> wrappedChild = new HierarchyWrapper<>(
-                            new TestBean(string + "/" + string2), rootBean,
-                            true);
-                    itemToWrapperMap.put(childBean, wrappedChild);
-                    wrapperToItemMap.put(wrappedChild, childBean);
-                    return childBean;
-                }).collect(Collectors.toList());
-
-                wrappedParent.setChildren(new LinkedHashSet<>(children));
-
-                rootNodes.put(rootBean, wrappedParent);
-            });
+        public HierarchicalTestBean(String parentId, int depth, int index) {
+            id = (parentId == null ? "" : parentId) + "/" + depth + "/" + index;
+            this.depth = depth;
+            this.index = index;
         }
 
-        @Override
-        public int getDepth(TestBean item) {
-            int depth = 0;
-            while (getItem(item) != null) {
-                item = getItem(item).getParent();
-                depth++;
-            }
+        public int getDepth() {
             return depth;
         }
 
-        @Override
-        public boolean isInMemory() {
-            return true;
-        }
-
-        @Override
-        public void refreshItem(TestBean item) {
-            // NO-OP
-        }
-
-        @Override
-        public void refreshAll() {
-            // NO-OP
-        }
-
-        @Override
-        public Registration addDataProviderListener(
-                DataProviderListener<TestBean> listener) {
-            return () -> {
-            };
-        }
-
-        private List<TestBean> getAllItems() {
-            return new ArrayList<>(itemToWrapperMap.keySet());
-        }
-
-        private List<TestBean> getVisibleItemsRecursive(
-                Collection<HierarchyWrapper<TestBean>> wrappedItems) {
-            List<TestBean> items = new ArrayList<>();
-
-            wrappedItems.forEach(wrappedItem -> {
-                items.add(wrapperToItemMap.get(wrappedItem));
-                if (!wrappedItem.isCollapsed()) {
-                    List<HierarchyWrapper<TestBean>> wrappedChildren = wrappedItem
-                            .getChildren().stream()
-                            .map(childItem -> getItem(childItem))
-                            .collect(Collectors.toList());
-                    items.addAll(getVisibleItemsRecursive(wrappedChildren));
-                }
-            });
-            return items;
+        public int getIndex() {
+            return index;
         }
 
-        @Override
-        public int size(Query<TestBean, Void> query) {
-            return getVisibleItemsRecursive(rootNodes.values()).size();
-        }
-
-        @Override
-        public Stream<TestBean> fetch(Query<TestBean, Void> query) {
-            return getVisibleItemsRecursive(rootNodes.values()).stream();
-        }
-
-        @Override
-        public boolean isRoot(TestBean item) {
-            return getItem(item).getParent() == null;
+        public String getId() {
+            return id;
         }
 
         @Override
-        public TestBean getParent(TestBean item) {
-            return getItem(item).getParent();
+        public String toString() {
+            return depth + " | " + index;
         }
 
         @Override
-        public boolean isCollapsed(TestBean item) {
-            return getItem(item).isCollapsed();
+        public int hashCode() {
+            final int prime = 31;
+            int result = 1;
+            result = prime * result + ((id == null) ? 0 : id.hashCode());
+            return result;
         }
 
         @Override
-        public boolean hasChildren(TestBean item) {
-            return !getItem(item).getChildren().isEmpty();
-        }
-
-        @Override
-        public void setCollapsed(TestBean item, boolean b) {
-            getItem(item).setCollapsed(b);
+        public boolean equals(Object obj) {
+            if (this == obj) {
+                return true;
+            }
+            if (obj == null) {
+                return false;
+            }
+            if (getClass() != obj.getClass()) {
+                return false;
+            }
+            HierarchicalTestBean other = (HierarchicalTestBean) obj;
+            if (id == null) {
+                if (other.id != null) {
+                    return false;
+                }
+            } else if (!id.equals(other.id)) {
+                return false;
+            }
+            return true;
         }
 
-        private HierarchyWrapper<TestBean> getItem(TestBean item) {
-            return itemToWrapperMap.get(item);
-        }
     }
 }

uitest/src/main/java/com/vaadin/tests/components/treegrid/TreeGridScrolling.java

@@ -0,0 +1,42 @@
+package com.vaadin.tests.components.treegrid;
+
+import com.vaadin.annotations.Widgetset;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.tests.components.treegrid.TreeGridBasicFeatures.HierarchicalTestBean;
+import com.vaadin.ui.TreeGrid;
+
+@Widgetset("com.vaadin.DefaultWidgetSet")
+public class TreeGridScrolling extends AbstractTestUI {
+
+    public static final int DEFAULT_NODES = 20;
+    public static final int DEFAULT_DEPTH = 3;
+    public static final String NODES_PARAMETER = "nodes";
+    public static final String DEPTH_PARAMETER = "depth";
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        int depth = DEFAULT_DEPTH;
+        if (request.getParameter(DEPTH_PARAMETER) != null) {
+            depth = Integer.parseInt(request.getParameter(DEPTH_PARAMETER));
+        }
+        int nodes = DEFAULT_NODES;
+        if (request.getParameter(NODES_PARAMETER) != null) {
+            nodes = Integer.parseInt(request.getParameter(NODES_PARAMETER));
+        }
+
+        TreeGrid<HierarchicalTestBean> grid = new TreeGrid<>();
+        grid.setSizeFull();
+        grid.addColumn(HierarchicalTestBean::toString).setCaption("String")
+                .setId("string");
+        grid.addColumn(HierarchicalTestBean::getDepth).setCaption("Depth")
+                .setId(DEPTH_PARAMETER);
+        grid.addColumn(HierarchicalTestBean::getIndex)
+                .setCaption("Index on this depth").setId("index");
+        grid.setHierarchyColumn("string");
+        grid.setDataProvider(new LazyHierarchicalDataProvider(nodes, depth));
+
+        addComponent(grid);
+    }
+
+}

uitest/src/main/java/com/vaadin/tests/components/treetable/TreeTableCacheOnPartialUpdates.java

@@ -62,7 +62,7 @@ public class TreeTableCacheOnPartialUpdates extends TestBase {
 
         @Override
         public String toString() {
-            return "TestBean [col1=" + col1 + ", col2=" + col2 + "]";
+            return "HierarchicalTestBean [col1=" + col1 + ", col2=" + col2 + "]";
         }
 
     }

uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridBasicFeaturesTest.java

@@ -1,19 +1,38 @@
 package com.vaadin.tests.components.treegrid;
 
+import java.util.Arrays;
+import java.util.Collection;
+
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized.Parameters;
 import org.openqa.selenium.Keys;
 import org.openqa.selenium.interactions.Actions;
 
 import com.vaadin.testbench.By;
 import com.vaadin.testbench.elements.TreeGridElement;
 import com.vaadin.tests.tb3.MultiBrowserTest;
+import com.vaadin.tests.tb3.ParameterizedTB3Runner;
 
+@RunWith(ParameterizedTB3Runner.class)
 public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
 
     private TreeGridElement grid;
 
+    public void setDataProvider(String dataProviderString) {
+        selectMenuPath("Component", "Features", "Set data provider",
+                dataProviderString);
+    }
+
+    @Parameters
+    public static Collection<String> getDataProviders() {
+        return Arrays.asList("LazyHierarchicalDataProvider",
+                "InMemoryHierarchicalDataProvider");
+    }
+
     @Before
     public void before() {
         openTestURL("theme=valo");
@@ -21,56 +40,59 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
     }
 
     @Test
+    @Ignore // currently no implementation exists for toggling from the server
+            // side
     public void toggle_collapse_server_side() {
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
-        selectMenuPath("Component", "Features", "Toggle expand", "a");
+        selectMenuPath("Component", "Features", "Toggle expand", "0 | 0");
         Assert.assertEquals(6, grid.getRowCount());
-        assertCellTexts(1, 0, new String[] { "a/a", "a/b", "a/c" });
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
 
-        selectMenuPath("Component", "Features", "Toggle expand", "a");
+        selectMenuPath("Component", "Features", "Toggle expand", "0 | 0");
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
         // collapsing a leaf should have no effect
-        selectMenuPath("Component", "Features", "Toggle expand", "a/a");
+        selectMenuPath("Component", "Features", "Toggle expand", "1 | 0");
         Assert.assertEquals(3, grid.getRowCount());
     }
 
     @Test
     public void non_leaf_collapse_on_click() {
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
 
-        // click the expander corresponding to "a"
+        // Should expand "0 | 0"
         grid.getRow(0).getCell(0)
                 .findElement(By.className("v-tree-grid-expander")).click();
         Assert.assertEquals(6, grid.getRowCount());
-        assertCellTexts(1, 0, new String[] { "a/a", "a/b", "a/c" });
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
 
-        // click the expander corresponding to "a"
+        // Should collapse "0 | 0"
         grid.getRow(0).getCell(0)
                 .findElement(By.className("v-tree-grid-expander")).click();
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
     }
 
     @Test
+    @Ignore // FIXME: remove ignore annotation once #8758 is done
     public void keyboard_navigation() {
         grid.getRow(0).getCell(0).click();
 
-        // Should expand "a"
+        // Should expand "0 | 0"
         new Actions(getDriver()).keyDown(Keys.ALT).sendKeys(Keys.RIGHT)
                 .keyUp(Keys.ALT).perform();
         Assert.assertEquals(6, grid.getRowCount());
-        assertCellTexts(1, 0, new String[] { "a/a", "a/b", "a/c" });
+        assertCellTexts(1, 0, new String[] { "1 | 0", "1 | 1", "1 | 2" });
 
-        // Should collapse "a"
+        // Should collapse "0 | 0"
         new Actions(getDriver()).keyDown(Keys.ALT).sendKeys(Keys.LEFT)
                 .keyUp(Keys.ALT).perform();
         Assert.assertEquals(3, grid.getRowCount());
-        assertCellTexts(0, 0, new String[] { "a", "b", "c" });
+        assertCellTexts(0, 0, new String[] { "0 | 0", "0 | 1", "0 | 2" });
     }
 
     @Test
@@ -81,7 +103,7 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
                 .isElementPresent(By.className("v-tree-grid-expander")));
 
         selectMenuPath("Component", "Features", "Set hierarchy column",
-                "Second column");
+                "depth");
 
         Assert.assertFalse(grid.getRow(0).getCell(0)
                 .isElementPresent(By.className("v-tree-grid-expander")));
@@ -89,7 +111,7 @@ public class TreeGridBasicFeaturesTest extends MultiBrowserTest {
                 .isElementPresent(By.className("v-tree-grid-expander")));
 
         selectMenuPath("Component", "Features", "Set hierarchy column",
-                "First column");
+                "string");
 
         Assert.assertTrue(grid.getRow(0).getCell(0)
                 .isElementPresent(By.className("v-tree-grid-expander")));

uitest/src/test/java/com/vaadin/tests/components/treegrid/TreeGridScrollingTest.java

@@ -0,0 +1,163 @@
+package com.vaadin.tests.components.treegrid;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import com.vaadin.testbench.elements.TreeGridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class TreeGridScrollingTest extends SingleBrowserTest {
+
+    @Test
+    public void testScrollingTree_expandCollapseFromBeginning_correctItemsShown() {
+        // TODO refactor this test to verify each row against a model, e.g. a
+        // InMemoryDataProvider, or the used lazy hierarchical data provider
+        openTestURL();
+
+        TreeGridElement grid = $(TreeGridElement.class).first();
+
+        Assert.assertEquals(grid.getRowCount(),
+                TreeGridScrolling.DEFAULT_NODES);
+
+        verifyRow(0, 0, 0);
+        verifyRow(10, 0, 10);
+        verifyRow(19, 0, 19);
+        verifyRow(10, 0, 10);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(0);
+
+        verifyRow(0, 0, 0);
+        verifyRow(1, 1, 0);
+        verifyRow(11, 1, 10);
+        verifyRow(20, 1, 19);
+        verifyRow(39, 0, 19);
+
+        // verifying in reverse order causes scrolling up
+        verifyRow(20, 1, 19);
+        verifyRow(11, 1, 10);
+        verifyRow(1, 1, 0);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(3);
+
+        verifyRow(0, 0, 0);
+
+        verifyRow(1, 1, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 1, 2);
+
+        verifyRow(4, 2, 0);
+
+        verifyRow(14, 2, 10);
+        verifyRow(23, 2, 19);
+        verifyRow(24, 1, 3);
+        verifyRow(40, 1, 19);
+        verifyRow(59, 0, 19);
+
+        // scroll back up
+
+        verifyRow(40, 1, 19);
+        verifyRow(24, 1, 3);
+        verifyRow(23, 2, 19);
+        verifyRow(14, 2, 10);
+
+        verifyRow(4, 2, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 1, 2);
+        verifyRow(1, 1, 0);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(2);
+
+        verifyRow(0, 0, 0);
+
+        verifyRow(1, 1, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 2, 0);
+        verifyRow(22, 2, 19);
+
+        verifyRow(23, 1, 2);
+        verifyRow(24, 2, 0);
+
+        verifyRow(43, 2, 19);
+        verifyRow(44, 1, 3);
+        verifyRow(60, 1, 19);
+        verifyRow(79, 0, 19);
+
+        // scroll back up
+        verifyRow(60, 1, 19);
+        verifyRow(44, 1, 3);
+        verifyRow(43, 2, 19);
+
+        verifyRow(24, 2, 0);
+        verifyRow(23, 1, 2);
+
+        verifyRow(22, 2, 19);
+        verifyRow(3, 2, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(1, 1, 0);
+
+        verifyRow(0, 0, 0);
+
+        grid.collapseWithClick(2);
+
+        verifyRow(0, 0, 0);
+
+        verifyRow(1, 1, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 1, 2);
+
+        verifyRow(4, 2, 0);
+
+        verifyRow(14, 2, 10);
+        verifyRow(23, 2, 19);
+        verifyRow(24, 1, 3);
+        verifyRow(40, 1, 19);
+        verifyRow(59, 0, 19);
+
+        // scroll back up
+
+        verifyRow(40, 1, 19);
+        verifyRow(24, 1, 3);
+        verifyRow(23, 2, 19);
+        verifyRow(14, 2, 10);
+
+        verifyRow(4, 2, 0);
+        verifyRow(2, 1, 1);
+        verifyRow(3, 1, 2);
+        verifyRow(1, 1, 0);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(3);
+
+        verifyRow(0, 0, 0);
+        verifyRow(1, 1, 0);
+        verifyRow(11, 1, 10);
+        verifyRow(20, 1, 19);
+        verifyRow(39, 0, 19);
+
+        // scroll back up
+
+        verifyRow(20, 1, 19);
+        verifyRow(11, 1, 10);
+        verifyRow(1, 1, 0);
+        verifyRow(0, 0, 0);
+
+        grid.expandWithClick(0);
+
+        verifyRow(0, 0, 0);
+        verifyRow(10, 0, 10);
+        verifyRow(19, 0, 19);
+        verifyRow(10, 0, 10);
+        verifyRow(0, 0, 0);
+    }
+
+    private void verifyRow(int rowActualIndex, int depth, int levelIndex) {
+        TreeGridElement grid = $(TreeGridElement.class).first();
+
+        Assert.assertEquals("Invalid row at index " + rowActualIndex,
+                depth + " | " + levelIndex,
+                grid.getCell(rowActualIndex, 0).getText());
+    }
+}