From 895930d636a9db4d076af132cebce075fc7183c2 Mon Sep 17 00:00:00 2001
From: Eric Davies <iamed2@gmail.com>
Date: Fri, 14 Jun 2019 12:02:49 -0500
Subject: [PATCH] Fix union and union! of intervals

---
 src/interval.jl  | 44 ++++++++++++++++++++++++++++++++++----------
 test/interval.jl | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 10 deletions(-)

diff --git a/src/interval.jl b/src/interval.jl
index 727b6e94..39379c49 100644
--- a/src/interval.jl
+++ b/src/interval.jl
@@ -279,22 +279,24 @@ end
 
 # There is power in a union.
 """
-    union(intervals::AbstractVector{<:AbstractInterval})
+    union(intervals::Union{AbstractInterval, AbstractVector{<:AbstractInterval}}, itrs...)
 
-Flattens a vector of overlapping intervals into a new, smaller vector containing only
+Flattens sets of overlapping intervals into a new, smaller vector containing only
 non-overlapping intervals.
 """
-function Base.union(intervals::AbstractVector{<:AbstractInterval})
-    return union!(convert(Vector{AbstractInterval}, intervals))
+function Base.union(
+    interval::Union{AbstractInterval, AbstractVector{<:AbstractInterval}}, intervals...,
+)
+    return union!(AbstractInterval[interval; intervals...])
 end
 
-"""
-    union!(intervals::AbstractVector{<:Union{Interval, AbstractInterval}})
+function Base.union!(intervals::AbstractVector{T}) where T <: AbstractInterval
+    if !(T <: Interval) && T != AbstractInterval
+        throw(ArgumentError(
+            "Cannot `union!` array of intervals of type $T as the type may change"
+        ))
+    end
 
-Flattens a vector of overlapping intervals in-place to be a smaller vector containing only
-non-overlapping intervals.
-"""
-function Base.union!(intervals::Union{AbstractVector{<:Interval}, AbstractVector{AbstractInterval}})
     sort!(intervals)
 
     i = 2
@@ -319,6 +321,28 @@ function Base.union!(intervals::Union{AbstractVector{<:Interval}, AbstractVector
     return intervals
 end
 
+"""
+    union!(intervals::AbstractVector{<:Union{Interval, AbstractInterval}}, itrs...)
+
+Flattens sets of overlapping intervals in-place into the first argument, making it a smaller
+vector containing only non-overlapping intervals.
+"""
+function union!(
+    intervals::AbstractVector{<:AbstractInterval},
+    interval_varargs...
+)
+    # imitates the behaviour of vcat/promote_eltypeof
+    for interval_vararg in interval_varargs
+        if interval_vararg isa AbstractInterval
+            push!(intervals, interval_vararg)
+        else
+            append!(intervals, interval_vararg)
+        end
+    end
+
+    return union!(intervals)
+end
+
 """
     superset(intervals::AbstractArray{<:AbstractInterval}) -> Interval
 
diff --git a/test/interval.jl b/test/interval.jl
index 7d9818bb..3956f440 100644
--- a/test/interval.jl
+++ b/test/interval.jl
@@ -536,6 +536,18 @@
     end
 
     @testset "union" begin
+        a = Interval(-100, -1)
+        b = Interval(-3, 10)
+
+        @test union(a) == AbstractInterval[a]
+        @test union(b) == AbstractInterval[b]
+        @test union(a, b) == AbstractInterval[Interval(-100, 10)]
+        @test union(b, a) == AbstractInterval[Interval(-100, 10)]
+        @test union([a, b]) == AbstractInterval[Interval(-100, 10)]
+        @test union([b, a]) == AbstractInterval[Interval(-100, 10)]
+        @test union!([a, b]) == [Interval(-100, 10)]
+        @test union!([b, a]) == [Interval(-100, 10)]
+
         intervals = [
             Interval(-100, -1, Inclusivity(false, false)),
             Interval(-10, -1, Inclusivity(false, false)),
@@ -572,5 +584,39 @@
             Interval(-10, -1, Inclusivity(true, true))
         ]
         @test union(intervals) == [Interval(-100, -1, Inclusivity(false, true))]
+
+        # anchored
+        intervals = [
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 12, 7, tz"UTC"), Inclusivity(false, true)),
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 12, 7, tz"UTC"), Inclusivity(false, true)),
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 13, 7, tz"UTC"), Inclusivity(false, true)),
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 13, 7, tz"UTC"), Inclusivity(false, true)),
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 14, 7, tz"UTC"), Inclusivity(false, true)),
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 14, 7, tz"UTC"), Inclusivity(false, true)),
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 12, 8, tz"UTC"), Inclusivity(false, true)),
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 12, 8, tz"UTC"), Inclusivity(false, true)),
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 13, 8, tz"UTC"), Inclusivity(false, true)),
+            HourEnding{ZonedDateTime}(ZonedDateTime(2013, 2, 13, 8, tz"UTC"), Inclusivity(false, true)),
+        ]
+
+        expected = [
+            Interval{ZonedDateTime}(ZonedDateTime(2013, 2, 12, 6, tz"UTC"), ZonedDateTime(2013, 2, 12, 8, tz"UTC"), Inclusivity(false, true)),
+            Interval{ZonedDateTime}(ZonedDateTime(2013, 2, 13, 6, tz"UTC"), ZonedDateTime(2013, 2, 13, 8, tz"UTC"), Inclusivity(false, true)),
+            Interval{ZonedDateTime}(ZonedDateTime(2013, 2, 14, 6, tz"UTC"), ZonedDateTime(2013, 2, 14, 7, tz"UTC"), Inclusivity(false, true)),
+        ]
+
+        @test union(intervals) == expected
+        @test_throws ArgumentError union!(intervals)
+
+        abstract_intervals = convert(Vector{AbstractInterval}, intervals)
+        @test union(abstract_intervals) == expected
+        @test length(abstract_intervals) == length(intervals)  # test that union doesn't mutate
+        @test union!(abstract_intervals) == expected
+        @test abstract_intervals == expected  # test that union! does mutate
+
+        # varargs
+        @test union(intervals...) == expected
+        @test_throws MethodError union!(intervals...)
+        @test union!(abstract_intervals[1:1], abstract_intervals[2:3], abstract_intervals[4:end]...) == expected
     end
 end