From e9df2a00468820866ba5741cd171abd5ad968ebd Mon Sep 17 00:00:00 2001
From: Evert van Brussel <ecvanbrussel@gmail.com>
Date: Sun, 13 May 2018 15:36:06 +0200
Subject: [PATCH] Bugfix and code clean up

---
 Sources/Gate/Gate.swift      | 36 ++++++++++++++++++++++--------------
 Sources/Gate/Policy.swift    | 12 +++++++++---
 Sources/Gate/TypeTuple.swift | 12 +++++++++++-
 3 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/Sources/Gate/Gate.swift b/Sources/Gate/Gate.swift
index d46298b..8440793 100644
--- a/Sources/Gate/Gate.swift
+++ b/Sources/Gate/Gate.swift
@@ -1,7 +1,11 @@
-public class Gate<Ability: AbilitySet> {
+public class Gate<Ability> where Ability: AbilitySet {
     public enum Mode {
         case giveRights
         case takeRights
+        
+        var defaultPolicy: Bool {
+            return self == .takeRights
+        }
     }
     
     let mode: Mode
@@ -43,33 +47,37 @@ extension Gate {
 }
 
 extension Gate {
-    public func check<User, Object>(_ user: User?, can ability: Ability, _ object: Object?) -> Bool {
-        if let generalAbilities = getAbilities(
+    public func check<User,Object>(_ user: User?, can ability: Ability, _ object: Object?) -> Bool {
+        let generalAbilities = getAbilities(
             from: policies[User.self, Any.self],
             user: user, object: object
-        ), !checkAllPolicies {
+        )
+        
+        if !checkAllPolicies, let generalAbilities = generalAbilities {
             return hasPermission(for: ability, given: generalAbilities)
         }
-
-        if let specificAbilities = getAbilities(
+        
+        let specificAbilities = getAbilities(
             from: policies[User.self, Object.self],
             user: user, object: object
-        ) {
-            return hasPermission(for: ability, given: specificAbilities)
+        )
+
+        if let specificAbilities = specificAbilities {
+            return hasPermission(for: ability, given: specificAbilities.union(generalAbilities ?? []))
         }
         
-        return mode == .takeRights
+        return mode.defaultPolicy
     }
     
-    public func check<User, Object>(_ user: User?, cannot ability: Ability, _ object: Object?) -> Bool {
+    public func check<User,Object>(_ user: User?, cannot ability: Ability, _ object: Object?) -> Bool {
         return !check(user, can: ability, object)
     }
     
-    public func check<User, Object>(_ user: User?, can ability: Ability, _ type: Object.Type) -> Bool {
+    public func check<User,Object>(_ user: User?, can ability: Ability, _ type: Object.Type) -> Bool {
         return check(user, can: ability, Object?.none)
     }
     
-    public func check<User, Object>(_ user: User?, cannot ability: Ability, _ type: Object.Type) -> Bool {
+    public func check<User,Object>(_ user: User?, cannot ability: Ability, _ type: Object.Type) -> Bool {
         return !check(user, can: ability, type)
     }
     
@@ -95,13 +103,13 @@ extension Gate {
 }
 
 extension Gate {
-    public func ensure<User, Object>(_ user: User?, can ability: Ability, _ object: Object?) throws {
+    public func ensure<User,Object>(_ user: User?, can ability: Ability, _ object: Object?) throws {
         guard check(user, can: ability, object) else {
             throw Unauthorized(user: user, ability: ability, object: object)
         }
     }
 
-    public func ensure<User, Object>(_ user: User?, can ability: Ability, _ type: Object.Type) throws {
+    public func ensure<User,Object>(_ user: User?, can ability: Ability, _ type: Object.Type) throws {
         try ensure(user, can: ability, Object?.none)
     }
 }
diff --git a/Sources/Gate/Policy.swift b/Sources/Gate/Policy.swift
index f4bb5df..7615204 100644
--- a/Sources/Gate/Policy.swift
+++ b/Sources/Gate/Policy.swift
@@ -1,10 +1,13 @@
-struct Policy<User, Object, Ability: AbilitySet> {
+struct Policy<User,Object,Ability> where Ability: AbilitySet {
     let getAbilities: (User?, Object?) -> Ability?
     
     init(_ policy: @escaping (User?, Object?) -> Ability?) {
         getAbilities = policy
     }
-    
+}
+
+// MARK: User only policies, AKA before policies
+extension Policy {
     init(_ policy: @escaping (User?) -> Ability?) {
         self.init { user, _ in policy(user) }
     }
@@ -15,7 +18,10 @@ struct Policy<User, Object, Ability: AbilitySet> {
             return policy(user)
         }
     }
-    
+}
+
+// MARK: Policies with a non-optional user and/or object
+extension Policy {
     init(_ policy: @escaping (User?, Object) -> Ability?) {
         self.init { user, object in
             guard let object = object else { return nil }
diff --git a/Sources/Gate/TypeTuple.swift b/Sources/Gate/TypeTuple.swift
index 4dd0a70..dfc4cf9 100644
--- a/Sources/Gate/TypeTuple.swift
+++ b/Sources/Gate/TypeTuple.swift
@@ -7,7 +7,17 @@ struct TypeTuple: Hashable {
 }
 
 extension Dictionary where Key == TypeTuple, Value == [Any] {
-    subscript<User,Object,Ability:AbilitySet>(_ userType: User.Type,_ objectType: Object.Type) -> [Policy<User,Object,Ability>] {
+    subscript<User,Object,Ability>(_ userType: User.Type, _ objectType: Object.Type, _: Ability.Type) -> [Policy<User,Object,Ability>] {
+        get {
+            return self[userType, objectType]
+        }
+        
+        set(newValue) {
+            self[userType, objectType] = newValue
+        }
+    }
+    
+    subscript<User,Object,Ability>(_ userType: User.Type, _ objectType: Object.Type) -> [Policy<User,Object,Ability>] {
         get {
             let hash = TypeTuple(User.self, Object.self)
             return self[hash] as? [Policy<User,Object,Ability>] ?? []