From 63fe0f09656b5c8f01f9c6113243ddaea9caeb97 Mon Sep 17 00:00:00 2001
From: Robert Maynard <robert.maynard@kitware.com>
Date: Wed, 15 May 2019 11:42:02 -0400
Subject: [PATCH] DispatcherReduceByKey uses the scheduling_range like all
 other dispatchers

---
 vtkm/worklet/DispatcherReduceByKey.h   |  4 ++--
 vtkm/worklet/internal/DispatcherBase.h | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/vtkm/worklet/DispatcherReduceByKey.h b/vtkm/worklet/DispatcherReduceByKey.h
index b6a0859cd..bc7b0f4c0 100644
--- a/vtkm/worklet/DispatcherReduceByKey.h
+++ b/vtkm/worklet/DispatcherReduceByKey.h
@@ -59,8 +59,8 @@ public:
     const InputDomainType& inputDomain = invocation.GetInputDomain();
 
     // Now that we have the input domain, we can extract the range of the
-    // scheduling and call BadicInvoke.
-    this->BasicInvoke(invocation, inputDomain.GetInputRange());
+    // scheduling and call BasicInvoke.
+    this->BasicInvoke(invocation, internal::scheduling_range(inputDomain));
   }
 };
 }
diff --git a/vtkm/worklet/internal/DispatcherBase.h b/vtkm/worklet/internal/DispatcherBase.h
index 5310f5c68..e8cf5fc7b 100644
--- a/vtkm/worklet/internal/DispatcherBase.h
+++ b/vtkm/worklet/internal/DispatcherBase.h
@@ -48,6 +48,8 @@ namespace vtkm
 {
 namespace worklet
 {
+template <typename T>
+class VTKM_ALWAYS_EXPORT Keys;
 namespace internal
 {
 
@@ -57,6 +59,13 @@ inline auto scheduling_range(const Domain& inputDomain) -> decltype(inputDomain.
   return inputDomain.GetNumberOfValues();
 }
 
+template <typename KeyType>
+inline auto scheduling_range(const vtkm::worklet::Keys<KeyType>& inputDomain)
+  -> decltype(inputDomain.GetInputRange())
+{
+  return inputDomain.GetInputRange();
+}
+
 template <typename Domain>
 inline auto scheduling_range(const Domain* const inputDomain)
   -> decltype(inputDomain->GetNumberOfValues())
@@ -64,6 +73,13 @@ inline auto scheduling_range(const Domain* const inputDomain)
   return inputDomain->GetNumberOfValues();
 }
 
+template <typename KeyType>
+inline auto scheduling_range(const vtkm::worklet::Keys<KeyType>* const inputDomain)
+  -> decltype(inputDomain->GetInputRange())
+{
+  return inputDomain->GetInputRange();
+}
+
 template <typename Domain, typename SchedulingRangeType>
 inline auto scheduling_range(const Domain& inputDomain, SchedulingRangeType type)
   -> decltype(inputDomain.GetSchedulingRange(type))