2 files changed, 39 insertions, 3 deletions

diff --git a/stats/statsview.cpp b/stats/statsview.cpp
index 09a984bdb..148611cfc 100644
--- a/stats/statsview.cpp
+++ b/stats/statsview.cpp
@@ -39,6 +39,7 @@ StatsView::StatsView(QQuickItem *parent) : QQuickItem(parent),
 	xAxis(nullptr),
 	yAxis(nullptr),
 	draggedItem(nullptr),
+	restrictDives(false),
 	rootNode(nullptr)
 {
 	setFlag(ItemHasContents, true);
@@ -496,6 +497,25 @@ void StatsView::reset()
 	grid.reset();
 }
 
+void StatsView::restrictToSelection()
+{
+	restrictedDives = getDiveSelection();
+	std::sort(restrictedDives.begin(), restrictedDives.end()); // Sort by pointer for quick lookup
+	restrictDives = true;
+	plot(state);
+}
+
+void StatsView::unrestrict()
+{
+	restrictDives = false;
+	plot(state);
+}
+
+int StatsView::restrictionCount() const
+{
+	return restrictDives ? (int)restrictedDives.size() : -1;
+}
+
 void StatsView::plot(const StatsState &stateIn)
 {
 	state = stateIn;
@@ -518,7 +538,19 @@ void StatsView::plotChart()
 		return;
 	reset();
 
-	const std::vector<dive *> dives = DiveFilter::instance()->visibleDives();
+	std::vector<dive *> dives;
+	if (restrictDives) {
+		std::vector<dive *> visible = DiveFilter::instance()->visibleDives();
+		dives.reserve(visible.size());
+		for (dive *d: visible) {
+			// binary search
+			auto it = std::lower_bound(restrictedDives.begin(), restrictedDives.end(), d);
+			if (it != restrictedDives.end() && *it == d)
+				dives.push_back(d);
+		}
+	} else {
+		dives = DiveFilter::instance()->visibleDives();
+	}
 	switch (state.type) {
 	case ChartType::DiscreteBar:
 		return plotBarChart(dives, state.subtype, state.var1, state.var1Binner, state.var2,
diff --git a/stats/statsview.h b/stats/statsview.h
index 6e0fe8967..204d4f568 100644
--- a/stats/statsview.h
+++ b/stats/statsview.h
@@ -46,13 +46,16 @@ public:
 
 	void plot(const StatsState &state);
 	void updateFeatures(const StatsState &state);	// Updates the visibility of chart features, such as legend, regression, etc.
+	void restrictToSelection();
+	void unrestrict();
+	int restrictionCount() const;			// <0: no restriction
 	QQuickWindow *w() const;			// Make window available to items
 	QSizeF size() const;
 	QRectF plotArea() const;
 	void addQSGNode(QSGNode *node, ChartZValue z);	// Must only be called in render thread!
 	void registerChartItem(ChartItem &item);
 	void registerDirtyChartItem(ChartItem &item);
-	void emergencyShutdown();			// Called when QQuick decides to delete out root node.
+	void emergencyShutdown();			// Called when QQuick decides to delete our root node.
 
 	// Create a chart item and add it to the scene.
 	// The item must not be deleted by the caller, but can be
@@ -131,7 +134,6 @@ private:
 	// Helper functions to add feature to the chart
 	void addLineMarker(double pos, double low, double high, const QPen &pen, bool isHorizontal);
 
-
 	StatsState state;
 	QFont titleFont;
 	std::vector<std::unique_ptr<StatsSeries>> series;
@@ -148,6 +150,8 @@ private:
 	QPointF dragStartMouse, dragStartItem;
 	SelectionModifier selectionModifier;
 	std::vector<dive *> oldSelection;
+	bool restrictDives;
+	std::vector<dive *> restrictedDives;	// sorted by pointer for quick lookup.
 
 	void hoverEnterEvent(QHoverEvent *event) override;
 	void hoverMoveEvent(QHoverEvent *event) override;