Logical checkpoint of ModelDownloader implementation (#7259)

* Restores ML Pods after M77.

* Fix Package.swift

* Re-add catalyst to GHA workflow.

* Includes the implementation of: 1) download types, 2) list models, 3) delete model.

* Add protocol to handle writing to user defaults.

* Add protocol to handle writing to user defaults.

* Add empty init to ModelDownloadConditions.

* Copyright year + other minor fixes.

* TODO to use FirebaseApp internal init for test app.

* Test if download progress is in range.

* Remove not nil test for progress.

* Remove not nil test for progress.

Author: manjanac

FirebaseMLModelDownloader/Sources/CustomModel.swift

@@ -1,4 +1,4 @@
-// Copyright 2020 Google LLC
+// Copyright 2021 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

FirebaseMLModelDownloader/Sources/DownloaderUserDefaults.swift

@@ -0,0 +1,20 @@
+// Copyright 2021 Google LLC
+//
+// 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 to write info to user defaults.
+protocol DownloaderUserDefaultsWriteable {
+  func writeToDefaults(_ defaults: UserDefaults, appName: String)
+}

FirebaseMLModelDownloader/Sources/LocalModelInfo.swift

@@ -1,4 +1,4 @@
-// Copyright 2020 Google LLC
+// Copyright 2021 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -33,12 +33,6 @@ class LocalModelInfo {
   /// Local path of the model.
   let path: String
 
-  /// Get user defaults key prefix.
-  private static func getUserDefaultsKeyPrefix(appName: String, modelName: String) -> String {
-    let bundleID = Bundle.main.bundleIdentifier ?? ""
-    return "\(bundleID).\(appName).\(modelName)"
-  }
-
   init(name: String, downloadURL: URL, modelHash: String, size: Int, path: String) {
     self.name = name
     self.downloadURL = downloadURL
@@ -71,6 +65,15 @@ class LocalModelInfo {
     }
     self.init(name: name, downloadURL: url, modelHash: modelHash, size: size, path: path)
   }
+}
+
+/// Extension to write local model info to user defaults.
+extension LocalModelInfo: DownloaderUserDefaultsWriteable {
+  /// Get user defaults key prefix.
+  private static func getUserDefaultsKeyPrefix(appName: String, modelName: String) -> String {
+    let bundleID = Bundle.main.bundleIdentifier ?? ""
+    return "\(bundleID).\(appName).\(modelName)"
+  }
 
   /// Write local model info to user defaults.
   func writeToDefaults(_ defaults: UserDefaults, appName: String) {

FirebaseMLModelDownloader/Sources/ModelDownloadConditions.swift

@@ -1,4 +1,4 @@
-// Copyright 2020 Google LLC
+// Copyright 2021 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -15,4 +15,8 @@
 import Foundation
 
 /// Model download conditions.
-public struct ModelDownloadConditions {}
+// TODO: Implement model download conditions.
+public struct ModelDownloadConditions {
+  // TODO: Intentionally left blank until implementation.
+  public init() {}
+}

FirebaseMLModelDownloader/Sources/ModelDownloadTask.swift

@@ -1,4 +1,4 @@
-// Copyright 2020 Google LLC
+// Copyright 2021 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -84,15 +84,21 @@ extension ModelDownloadTask: URLSessionDownloadDelegate {
     return "fbml_model__\(appName)__\(remoteModelInfo.name)"
   }
 
+  func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
+    // TODO: Handle model download url expiry and other download errors
+  }
+
   func urlSession(_ session: URLSession,
                   downloadTask: URLSessionDownloadTask,
                   didFinishDownloadingTo location: URL) {
     assert(downloadTask == self.downloadTask)
     downloadStatus = .completed
-    let savedURL = ModelFileManager.modelsDirectory
-      .appendingPathComponent(downloadedModelFileName)
+    let modelFileURL = ModelFileManager.getDownloadedModelFilePath(
+      appName: appName,
+      modelName: remoteModelInfo.name
+    )
     do {
-      try ModelFileManager.moveFile(at: location, to: savedURL)
+      try ModelFileManager.moveFile(at: location, to: modelFileURL)
     } catch let error as DownloadError {
       DispatchQueue.main.async {
         self.downloadHandlers
@@ -106,7 +112,7 @@ extension ModelDownloadTask: URLSessionDownloadDelegate {
     }
 
     /// Generate local model info.
-    let localModelInfo = LocalModelInfo(from: remoteModelInfo, path: savedURL.absoluteString)
+    let localModelInfo = LocalModelInfo(from: remoteModelInfo, path: modelFileURL.absoluteString)
     /// Write model to user defaults.
     localModelInfo.writeToDefaults(defaults, appName: appName)
     /// Build model from model info.

FirebaseMLModelDownloader/Sources/ModelDownloader.swift

@@ -1,4 +1,4 @@
-// Copyright 2020 Google LLC
+// Copyright 2021 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -14,6 +14,7 @@
 
 import Foundation
 import FirebaseCore
+import FirebaseInstallations
 
 /// Possible errors with model downloading.
 public enum DownloadError: Error, Equatable {
@@ -37,6 +38,8 @@ public enum DownloadedModelError: Error, Equatable {
   case fileIOError(description: String)
   /// Model not found on device.
   case notFound
+  /// Other errors with description.
+  case internalError(description: String)
 }
 
 /// Possible ways to get a custom model.
@@ -53,14 +56,24 @@ public enum ModelDownloadType {
 public class ModelDownloader {
   /// Name of the app associated with this instance of ModelDownloader.
   private let appName: String
+  /// Current Firebase app options.
+  private let options: FirebaseOptions
+  /// Installations instance for current Firebase app.
+  private let installations: Installations
+  /// User defaults for model info.
+  private let userDefaults: UserDefaults
 
   /// Shared dictionary mapping app name to a specific instance of model downloader.
   // TODO: Switch to using Firebase components.
   private static var modelDownloaderDictionary: [String: ModelDownloader] = [:]
 
   /// Private init for downloader.
-  private init(app: FirebaseApp) {
+  private init(app: FirebaseApp, defaults: UserDefaults = .firebaseMLDefaults) {
     appName = app.name
+    options = app.options
+    installations = Installations.installations(app: app)
+    userDefaults = defaults
+
     NotificationCenter.default.addObserver(
       self,
       selector: #selector(deleteModelDownloader),
@@ -103,36 +116,202 @@ public class ModelDownloader {
                        conditions: ModelDownloadConditions,
                        progressHandler: ((Float) -> Void)? = nil,
                        completion: @escaping (Result<CustomModel, DownloadError>) -> Void) {
-    // TODO: Model download
-    let modelSize = Int()
-    let modelPath = String()
-    let modelHash = String()
+    switch downloadType {
+    case .localModel:
+      if let localModel = getLocalModel(modelName: modelName) {
+        DispatchQueue.main.async {
+          completion(.success(localModel))
+        }
+      } else {
+        getRemoteModel(
+          modelName: modelName,
+          progressHandler: progressHandler,
+          completion: completion
+        )
+      }
+    case .localModelUpdateInBackground:
+      if let localModel = getLocalModel(modelName: modelName) {
+        DispatchQueue.main.async {
+          completion(.success(localModel))
+        }
+        DispatchQueue.global(qos: .utility).async { [weak self] in
+          self?.getRemoteModel(
+            modelName: modelName,
+            progressHandler: nil,
+            completion: { result in
+              switch result {
+              // TODO: Log outcome of background download
+              case .success: break
+              case .failure: break
+              }
+            }
+          )
+        }
+      } else {
+        getRemoteModel(
+          modelName: modelName,
+          progressHandler: progressHandler,
+          completion: completion
+        )
+      }
 
-    let customModel = CustomModel(
-      name: modelName,
-      size: modelSize,
-      path: modelPath,
-      hash: modelHash
-    )
-    completion(.success(customModel))
-    completion(.failure(.notFound))
+    case .latestModel:
+      getRemoteModel(
+        modelName: modelName,
+        progressHandler: progressHandler,
+        completion: completion
+      )
+    }
   }
 
   /// Gets all downloaded models.
   public func listDownloadedModels(completion: @escaping (Result<Set<CustomModel>,
     DownloadedModelError>) -> Void) {
-    let customModels = Set<CustomModel>()
-    // TODO: List downloaded models
-    completion(.success(customModels))
-    completion(.failure(.notFound))
+    do {
+      let modelPaths = try ModelFileManager.contentsOfModelsDirectory()
+      var customModels = Set<CustomModel>()
+      for path in modelPaths {
+        guard let modelName = ModelFileManager.getModelNameFromFilePath(path) else {
+          completion(.failure(.internalError(description: "Invalid model file name.")))
+          return
+        }
+        guard let modelInfo = getLocalModelInfo(modelName: modelName) else {
+          completion(
+            .failure(.internalError(description: "Failed to get model info for model file."))
+          )
+          return
+        }
+        guard modelInfo.path == path.absoluteString else {
+          completion(
+            .failure(.internalError(description: "Outdated model paths in local storage."))
+          )
+          return
+        }
+        let model = CustomModel(localModelInfo: modelInfo)
+        customModels.insert(model)
+      }
+      completion(.success(customModels))
+    } catch let error as DownloadedModelError {
+      completion(.failure(error))
+    } catch {
+      completion(.failure(.internalError(description: error.localizedDescription)))
+    }
   }
 
   /// Deletes a custom model from device.
   public func deleteDownloadedModel(name modelName: String,
                                     completion: @escaping (Result<Void, DownloadedModelError>)
                                       -> Void) {
     // TODO: Delete previously downloaded model
-    completion(.success(()))
-    completion(.failure(.notFound))
+    guard let localModelInfo = getLocalModelInfo(modelName: modelName),
+      let localPath = URL(string: localModelInfo.path)
+    else {
+      completion(.failure(.notFound))
+      return
+    }
+    do {
+      try ModelFileManager.removeFile(at: localPath)
+      completion(.success(()))
+    } catch let error as DownloadedModelError {
+      completion(.failure(error))
+    } catch {
+      completion(.failure(.internalError(description: error.localizedDescription)))
+    }
+  }
+}
+
+extension ModelDownloader {
+  /// Return local model info only if the model info is available and the corresponding model file is already on device.
+  private func getLocalModelInfo(modelName: String) -> LocalModelInfo? {
+    guard let localModelInfo = LocalModelInfo(
+      fromDefaults: userDefaults,
+      name: modelName,
+      appName: appName
+    ) else {
+      return nil
+    }
+    /// There is local model info on device, but no model file at the expected path.
+    guard let localPath = URL(string: localModelInfo.path),
+      ModelFileManager.isFileReachable(at: localPath) else {
+      // TODO: Consider deleting local model info in user defaults.
+      return nil
+    }
+    return localModelInfo
+  }
+
+  /// Get model saved on device if available.
+  private func getLocalModel(modelName: String) -> CustomModel? {
+    guard let localModelInfo = getLocalModelInfo(modelName: modelName) else { return nil }
+    let model = CustomModel(localModelInfo: localModelInfo)
+    return model
+  }
+
+  /// Download and get model from server, unless the latest model is already available on device.
+  private func getRemoteModel(modelName: String,
+                              progressHandler: ((Float) -> Void)? = nil,
+                              completion: @escaping (Result<CustomModel, DownloadError>) -> Void) {
+    let localModelInfo = getLocalModelInfo(modelName: modelName)
+    let modelInfoRetriever = ModelInfoRetriever(
+      modelName: modelName,
+      options: options,
+      installations: installations,
+      appName: appName,
+      localModelInfo: localModelInfo
+    )
+    modelInfoRetriever.downloadModelInfo { result in
+      switch result {
+      case let .success(downloadModelInfoResult):
+        switch downloadModelInfoResult {
+        /// New model info was downloaded from server.
+        case let .modelInfo(remoteModelInfo):
+          let downloadTask = ModelDownloadTask(
+            remoteModelInfo: remoteModelInfo,
+            appName: self.appName,
+            defaults: self.userDefaults,
+            progressHandler: progressHandler,
+            completion: completion
+          )
+          downloadTask.resumeModelDownload()
+        /// Local model info is the latest model info.
+        case .notModified:
+          guard let localModel = self.getLocalModel(modelName: modelName) else {
+            DispatchQueue.main.async {
+              /// This can only happen if either local model info or the model file was suddenly wiped out in the middle of model info request and server response
+              // TODO: Consider handling: if model file is deleted after local model info is retrieved but before model info network call
+              completion(
+                .failure(
+                  .internalError(description: "Model unavailable due to deleted local model info.")
+                )
+              )
+            }
+            return
+          }
+
+          DispatchQueue.main.async {
+            completion(.success(localModel))
+          }
+        }
+      case let .failure(downloadError):
+        DispatchQueue.main.async {
+          completion(.failure(downloadError))
+        }
+      }
+    }
+  }
+}
+
+/// Model downloader extension for testing.
+extension ModelDownloader {
+  /// Model downloader instance for testing.
+  // TODO: Consider using protocols
+  static func modelDownloaderWithDefaults(_ defaults: UserDefaults,
+                                          app: FirebaseApp) -> ModelDownloader {
+    if let downloader = modelDownloaderDictionary[app.name] {
+      return downloader
+    } else {
+      let downloader = ModelDownloader(app: app, defaults: defaults)
+      modelDownloaderDictionary[app.name] = downloader
+      return downloader
+    }
   }
 }