implement experiment manager

Sources/BrowserServicesKit/PrivacyConfig/ExperimentCohortsManager.swift

@@ -0,0 +1,107 @@
+//
+//  ExperimentCohortsManager.swift
+//
+//  Copyright © 2024 DuckDuckGo. All rights reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+import Foundation
+
+struct ExperimentSubfeature {
+    let subfeatureID: SubfeatureID
+    let cohorts: [PrivacyConfigurationData.Cohort]
+}
+
+typealias CohortID = String
+typealias SubfeatureID = String
+
+struct ExperimentData: Codable, Equatable {
+    let cohort: String
+    let enrollmentDate: Date
+}
+
+typealias Experiments = [String: ExperimentData]
+
+protocol ExperimentCohortsManaging {
+    /// Retrieves the cohort ID associated with the specified subfeature.
+    /// - Parameter subfeatureID: The name of the experiment subfeature for which the cohort ID is needed.
+    /// - Returns: The cohort ID as a `String` if one exists; otherwise, returns `nil`.
+    func cohort(for subfeatureID: SubfeatureID) -> CohortID?
+
+    /// Retrieves the enrollment date for the specified subfeature.
+    /// - Parameter subfeatureID: The name of the experiment subfeature for which the enrollment date is needed.
+    /// - Returns: The `Date` of enrollment if one exists; otherwise, returns `nil`.
+    func enrollmentDate(for subfeatureID: SubfeatureID) -> Date?
+
+    /// Assigns a cohort to the given subfeature based on defined weights and saves it to UserDefaults.
+    /// - Parameter subfeature: The ExperimentSubfeature to which a cohort needs to be assigned to.
+    /// - Returns: The name of the assigned cohort, or `nil` if no cohort could be assigned.
+    func assignCohort(to subfeature: ExperimentSubfeature) -> CohortID?
+
+    /// Removes the assigned cohort data for the specified subfeature.
+    /// - Parameter subfeatureID: The name of the experiment subfeature for which the cohort data should be removed.
+    func removeCohort(from subfeatureID: SubfeatureID)
+}
+
+final class ExperimentCohortsManager: ExperimentCohortsManaging {
+
+    private var store: ExperimentsDataStoring
+    private let randomizer: (Range<Double>) -> Double
+
+    init(store: ExperimentsDataStoring = ExperimentsDataStore(), randomizer: @escaping (Range<Double>) -> Double) {
+        self.store = store
+        self.randomizer = randomizer
+    }
+
+    func cohort(for subfeatureID: SubfeatureID) -> CohortID? {
+        guard let experiments = store.experiments else { return nil }
+        return experiments[subfeatureID]?.cohort
+    }
+
+    func enrollmentDate(for subfeatureID: SubfeatureID) -> Date? {
+        guard let experiments = store.experiments else { return nil }
+        return experiments[subfeatureID]?.enrollmentDate
+    }
+
+    func assignCohort(to subfeature: ExperimentSubfeature) -> CohortID? {
+        let cohorts = subfeature.cohorts
+        let totalWeight = cohorts.map(\.weight).reduce(0, +)
+        guard totalWeight > 0 else { return nil }
+
+        let randomValue = randomizer(0..<Double(totalWeight))
+        var cumulativeWeight = 0.0
+
+        for cohort in cohorts {
+            cumulativeWeight += Double(cohort.weight)
+            if randomValue < cumulativeWeight {
+                saveCohort(cohort.name, in: subfeature.subfeatureID)
+                return cohort.name
+            }
+        }
+        return nil
+    }
+
+    func removeCohort(from subfeatureID: SubfeatureID) {
+        guard var experiments = store.experiments else { return }
+        experiments.removeValue(forKey: subfeatureID)
+        store.experiments = experiments
+    }
+
+    private func saveCohort(_ cohort: CohortID, in experimentID: SubfeatureID) {
+        var experiments = store.experiments ?? Experiments()
+        let experimentData = ExperimentData(cohort: cohort, enrollmentDate: Date())
+        experiments[experimentID] = experimentData
+        store.experiments = experiments
+    }
+}

Sources/BrowserServicesKit/PrivacyConfig/ExperimentsDataStore.swift

@@ -0,0 +1,58 @@
+//
+//  ExperimentsDataStore.swift
+//
+//  Copyright © 2024 DuckDuckGo. All rights reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+import Foundation
+
+protocol ExperimentsDataStoring {
+    var experiments: Experiments? { get set }
+}
+
+protocol LocalDataStoring {
+    func data(forKey defaultName: String) -> Data?
+    func set(_ value: Any?, forKey defaultName: String)
+}
+
+struct ExperimentsDataStore: ExperimentsDataStoring {
+
+    private enum Constants {
+        static let experimentsDataKey = "ExperimentsData"
+    }
+    private let localDataStoring: LocalDataStoring
+    private let decoder = JSONDecoder()
+    private let encoder = JSONEncoder()
+
+    init(localDataStoring: LocalDataStoring = UserDefaults.standard) {
+        self.localDataStoring = localDataStoring
+        encoder.dateEncodingStrategy = .secondsSince1970
+        decoder.dateDecodingStrategy = .secondsSince1970
+    }
+
+    var experiments: Experiments? {
+        get {
+            guard let savedData = localDataStoring.data(forKey: Constants.experimentsDataKey) else { return nil }
+            return try? decoder.decode(Experiments.self, from: savedData)
+        }
+        set {
+            if let encodedData = try? encoder.encode(newValue) {
+                localDataStoring.set(encodedData, forKey: Constants.experimentsDataKey)
+            }
+        }
+    }
+}
+
+extension UserDefaults: LocalDataStoring {}

Sources/BrowserServicesKit/PrivacyConfig/PrivacyConfigurationData.swift

@@ -36,6 +36,20 @@ public struct PrivacyConfigurationData {
         static public let enabled = "enabled"
     }
 
+    public struct Cohort {
+        public let name: String
+        public let weight: Int
+
+        public init?(json: [String: Any]) {
+            guard let name = json["name"] as? String,
+                  let weight = json["weight"] as? Int else {
+                return nil
+            }
+
+            self.name = name
+            self.weight = weight
+        }
+    }
     public let features: [FeatureName: PrivacyFeature]
     public let trackerAllowlist: TrackerAllowlist
     public let unprotectedTemporary: [ExceptionEntry]
@@ -121,6 +135,7 @@ public struct PrivacyConfigurationData {
                 case state
                 case minSupportedVersion
                 case rollout
+                case cohorts
             }
 
             public struct Rollout: Hashable {
@@ -157,6 +172,7 @@ public struct PrivacyConfigurationData {
             public let state: FeatureState
             public let minSupportedVersion: FeatureSupportedVersion?
             public let rollout: Rollout?
+            public let cohorts: [Cohort]?
 
             public init?(json: [String: Any]) {
                 guard let state = json[CodingKeys.state.rawValue] as? String else {
@@ -171,6 +187,13 @@ public struct PrivacyConfigurationData {
                 } else {
                     self.rollout = nil
                 }
+
+                if let cohortData = json[CodingKeys.cohorts.rawValue] as? [[String: Any]] {
+                    let parsedCohorts = cohortData.compactMap { Cohort(json: $0) }
+                    cohorts = parsedCohorts.isEmpty ? nil : parsedCohorts
+                } else {
+                    cohorts = nil
+                }
             }
         }