diff --git a/ArrayDiff/VMBArrayDiff.h b/ArrayDiff/VMBArrayDiff.h
index a03559c..ce6cb73 100644
--- a/ArrayDiff/VMBArrayDiff.h
+++ b/ArrayDiff/VMBArrayDiff.h
@@ -21,7 +21,7 @@ typedef NS_ENUM(NSInteger, VMBArrayDiffChangeType) {
 @end
 
 
-extern NSSet *VMBChangesByDiffingArrays(NSArray *left, NSArray *right) __attribute__((overloadable));
-extern NSSet *VMBChangesByDiffingArrays(NSArray *left, NSArray *right, __nullable VMBArrayElementComparator comparator) __attribute__((overloadable));
+extern NSSet<VMBArrayDiffChange *> *VMBChangesByDiffingArrays(NSArray *left, NSArray *right) __attribute__((overloadable));
+extern NSSet<VMBArrayDiffChange *> *VMBChangesByDiffingArrays(NSArray *left, NSArray *right, __nullable VMBArrayElementComparator comparator) __attribute__((overloadable));
 
 NS_ASSUME_NONNULL_END
diff --git a/ArrayDiffTests/ArrayDiffTests.m b/ArrayDiffTests/ArrayDiffTests.m
index cfdd74e..0cfc62b 100644
--- a/ArrayDiffTests/ArrayDiffTests.m
+++ b/ArrayDiffTests/ArrayDiffTests.m
@@ -28,13 +28,13 @@ - (void)testDiff {
         @"comment",
     ];
 
-    NSSet *expectedDiff = [NSSet setWithArray:@[
+    NSSet<VMBArrayDiffChange *> *expectedDiff = [NSSet setWithArray:@[
         [[VMBArrayDiffChange alloc] initWithChangeKind:VMBArrayDiffChangeTypeInsert index:0],
         [[VMBArrayDiffChange alloc] initWithChangeKind:VMBArrayDiffChangeTypeDelete index:4],
         [[VMBArrayDiffChange alloc] initWithChangeKind:VMBArrayDiffChangeTypeInsert index:2],
     ]];
 
-    NSSet *diff = VMBChangesByDiffingArrays(left, right);
+    NSSet<VMBArrayDiffChange *> *diff = VMBChangesByDiffingArrays(left, right);
 
     XCTAssertEqualObjects(diff, expectedDiff);
 }