From 71679dfd1626737081b86127e6c547e37c77923f Mon Sep 17 00:00:00 2001
From: Aleksi Hietanen <aleksi@vaadin.com>
Date: Thu, 16 Mar 2017 08:53:38 +0200
Subject: Hierarchical data (#8842)

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

Fixes #8611, Fixes #8620
---
 server/src/main/java/com/vaadin/data/HasItems.java |   4 +-
 .../main/java/com/vaadin/data/HierarchyData.java   | 263 ++++++++++++
 .../provider/AbstractHierarchicalDataProvider.java |  34 ++
 .../com/vaadin/data/provider/DataCommunicator.java | 135 ++++++-
 .../provider/HierarchicalDataCommunicator.java     | 390 ++++++++++++++++++
 .../data/provider/HierarchicalDataProvider.java    |  82 +++-
 .../vaadin/data/provider/HierarchicalQuery.java    |  93 +++++
 .../com/vaadin/data/provider/HierarchyMapper.java  | 445 +++++++++++++++++++++
 .../provider/InMemoryHierarchicalDataProvider.java | 235 +++++++++++
 .../com/vaadin/data/provider/ListDataProvider.java |   3 +-
 server/src/main/java/com/vaadin/ui/Grid.java       |  72 +++-
 server/src/main/java/com/vaadin/ui/TreeGrid.java   | 265 +++++++++---
 12 files changed, 1923 insertions(+), 98 deletions(-)
 create mode 100644 server/src/main/java/com/vaadin/data/HierarchyData.java
 create mode 100644 server/src/main/java/com/vaadin/data/provider/AbstractHierarchicalDataProvider.java
 create mode 100644 server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
 create mode 100644 server/src/main/java/com/vaadin/data/provider/HierarchicalQuery.java
 create mode 100644 server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
 create mode 100644 server/src/main/java/com/vaadin/data/provider/InMemoryHierarchicalDataProvider.java

(limited to 'server/src/main/java/com')

diff --git a/server/src/main/java/com/vaadin/data/HasItems.java b/server/src/main/java/com/vaadin/data/HasItems.java
index 18dd9fc99d..c5b5382c9e 100644
--- a/server/src/main/java/com/vaadin/data/HasItems.java
+++ b/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();
diff --git a/server/src/main/java/com/vaadin/data/HierarchyData.java b/server/src/main/java/com/vaadin/data/HierarchyData.java
new file mode 100644
index 0000000000..df5fbbbbf5
--- /dev/null
+++ b/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);
+    }
+}
diff --git a/server/src/main/java/com/vaadin/data/provider/AbstractHierarchicalDataProvider.java b/server/src/main/java/com/vaadin/data/provider/AbstractHierarchicalDataProvider.java
new file mode 100644
index 0000000000..4b9825dcf7
--- /dev/null
+++ b/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> {
+
+}
diff --git a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
index 9080ad3442..ffa461e19f 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/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();
         }
diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
new file mode 100644
index 0000000000..75e131df57
--- /dev/null
+++ b/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);
+    }
+
+}
diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataProvider.java b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataProvider.java
index ec54a3a138..8dabdce2cc 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataProvider.java
+++ b/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);
 }
diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchicalQuery.java b/server/src/main/java/com/vaadin/data/provider/HierarchicalQuery.java
new file mode 100644
index 0000000000..6671464b31
--- /dev/null
+++ 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);
+    }
+}
diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java b/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
new file mode 100644
index 0000000000..7f423b39c9
--- /dev/null
+++ b/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());
+        }
+    }
+
+}
diff --git a/server/src/main/java/com/vaadin/data/provider/InMemoryHierarchicalDataProvider.java b/server/src/main/java/com/vaadin/data/provider/InMemoryHierarchicalDataProvider.java
new file mode 100644
index 0000000000..b86c3186c6
--- /dev/null
+++ b/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);
+    }
+}
diff --git a/server/src/main/java/com/vaadin/data/provider/ListDataProvider.java b/server/src/main/java/com/vaadin/data/provider/ListDataProvider.java
index 8fc6d4a364..fba2f94d9d 100644
--- a/server/src/main/java/com/vaadin/data/provider/ListDataProvider.java
+++ b/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
diff --git a/server/src/main/java/com/vaadin/ui/Grid.java b/server/src/main/java/com/vaadin/ui/Grid.java
index 19222d8c85..175b7d1ae5 100644
--- a/server/src/main/java/com/vaadin/ui/Grid.java
+++ b/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));
diff --git a/server/src/main/java/com/vaadin/ui/TreeGrid.java b/server/src/main/java/com/vaadin/ui/TreeGrid.java
index 46e6c99b59..82c857ac86 100644
--- a/server/src/main/java/com/vaadin/ui/TreeGrid.java
+++ b/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
-- 
cgit v1.2.3