Context.swift

// Sources/protoc-gen-swift/Context.swift - Overall code generation handling
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See http://swift.org/LICENSE.txt for license information
// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
// -----------------------------------------------------------------------------
///
/// The 'Context' wraps the CodeGeneratorRequest provided by protoc.  As such,
/// it is the only place that actually has access to all of the information provided
/// by protoc.
///
/// Much of protoc-gen-swift is based around two basic idioms:
///   - Each descriptor object is wrapped by a 'generator' that provides
///     additional data storage and logic
///   - Most methods are invoked with a reference to the Context class so
///     they can look up information about remote types.  Note that this
///     reference is never stored locally.
///
// -----------------------------------------------------------------------------
import Foundation
import PluginLibrary
import SwiftProtobuf

/*
 * A tool for looking up information about various types within
 * the overall context.
 */

typealias CodeGeneratorRequest = Google_Protobuf_Compiler_CodeGeneratorRequest
typealias CodeGeneratorResponse = Google_Protobuf_Compiler_CodeGeneratorResponse

extension CodeGeneratorRequest {
  func getMessageForPath(path: String) -> Google_Protobuf_DescriptorProto? {
    for f in protoFile {
      if let m = f.getMessageForPath(path: path) {
        return m
      }
    }
    return nil
  }

  func getMessageNameForPath(path: String) -> String? {
    for f in protoFile {
      if let m = f.getMessageNameForPath(path: path) {
        return m
      }
    }
    return nil
  }

  func getEnumNameForPath(path: String) -> String? {
    for f in protoFile {
      if let m = f.getEnumNameForPath(path: path) {
        return m
      }
    }
    return nil
  }

  func getSwiftNameForEnumCase(path: String, caseName: String) -> String {
    for f in protoFile {
      if let m = f.getSwiftNameForEnumCase(path: path, caseName: caseName) {
        return m
      }
    }
    fatalError("Unable to locate Enum case \(caseName) in path \(path)")
  }
}

extension CodeGeneratorResponse {
  init(error: String) {
    self.init()
    self.error = error
  }
}

extension CodeGeneratorResponse.File {
  init(name: String, content: String) {
    self.init()
    self.name = name
    self.content = content
  }
}

class GeneratorOptions {
  enum OutputNaming : String {
    case FullPath
    case PathToUnderscores
    case DropPath
  }

  enum Visibility : String {
    case Internal
    case Public
  }

  private(set) var outputNaming: OutputNaming = .FullPath
  private(set) var visibility: Visibility = .Internal

  init(parameter: String?) throws {
    for pair in parseParameter(string:parameter) {
      switch pair.key {
      case "FileNaming":
        if let naming = OutputNaming(rawValue: pair.value) {
          outputNaming = naming
        } else {
          throw GenerationError.invalidParameterValue(name: pair.key,
                                                      value: pair.value)
        }
      case "Visibility":
        if let value = Visibility(rawValue: pair.value) {
          visibility = value
        } else {
          throw GenerationError.invalidParameterValue(name: pair.key,
                                                      value: pair.value)
        }
      default:
        throw GenerationError.unknownParameter(name: pair.key)
      }
    }
  }
}

extension GeneratorOptions {
  /// Returns a string snippet to insert for the visibility
  var visibilitySourceSnippet: String {
    switch visibility {
    case .Internal:
      return ""
    case .Public:
      return "public "
    }
  }
}

class Context {
  let request: CodeGeneratorRequest
  let options: GeneratorOptions

  private(set) var parent = [String:String]()
  private(set) var fileByProtoName = [String:Google_Protobuf_FileDescriptorProto]()
  private(set) var enumByProtoName = [String:Google_Protobuf_EnumDescriptorProto]()
  private(set) var messageByProtoName = [String:Google_Protobuf_DescriptorProto]()
  private(set) var protoNameIsGroup = Set<String>()

  func swiftNameForProtoName(protoName: String, appending: String? = nil, separator: String = ".") -> String {
    let p = parent[protoName]
    if let e = enumByProtoName[protoName] {
      return swiftNameForProtoName(protoName: p!, appending: e.name, separator: separator)
    } else if let m = messageByProtoName[protoName] {
      let baseName: String
      if protoNameIsGroup.contains(protoName) {
        // TODO: Find a way to actually get to this line of code.
        // Then fix it to be whatever it should be.
        // If it can't be reached, assert an error in this case.
        baseName = "XXGROUPXX_" + m.name + "_XXGROUPXX"
      } else {
        baseName = m.name
      }
      let name: String
      if let a = appending {
        name = baseName + separator + a
      } else {
        name = baseName
      }
      return swiftNameForProtoName(protoName: p!, appending: name, separator: separator)
    } else if let f = fileByProtoName[protoName] {
      return f.swiftPrefix + (appending ?? "")
    }
    return ""
  }

  func getMessageForPath(path: String) -> Google_Protobuf_DescriptorProto? {
    return request.getMessageForPath(path: path)
  }

  func getMessageNameForPath(path: String) -> String? {
    return request.getMessageNameForPath(path: path)
  }

  func getEnumNameForPath(path: String) -> String? {
    return request.getEnumNameForPath(path: path)
  }

  init(request: CodeGeneratorRequest) throws {
    self.request = request
    self.options = try GeneratorOptions(parameter: request.parameter)
    for fileProto in request.protoFile {
      populateFrom(fileProto: fileProto)
    }
  }

  func populateFrom(fileProto: Google_Protobuf_FileDescriptorProto) {
    let prefix: String
    let pkg = fileProto.package
    if !pkg.isEmpty {
      prefix = "." + pkg
    } else {
      prefix = ""
    }
    fileByProtoName[prefix] = fileProto
    for e in fileProto.enumType {
      populateFrom(enumProto: e, prefix: prefix)
    }
    for m in fileProto.messageType {
      populateFrom(messageProto: m, prefix: prefix)
    }
    for f in fileProto.extension_p {
      if f.type == .group {
        protoNameIsGroup.insert(f.typeName)
      }
    }
  }

  func populateFrom(enumProto: Google_Protobuf_EnumDescriptorProto, prefix: String) {
    let name = prefix + "." + enumProto.name
    enumByProtoName[name] = enumProto
    parent[name] = prefix
  }

  func populateFrom(messageProto: Google_Protobuf_DescriptorProto, prefix: String) {
    let name = prefix + "." + messageProto.name
    parent[name] = prefix
    messageByProtoName[name] = messageProto
    for f in messageProto.field {
      if f.type == .group {
        protoNameIsGroup.insert(f.typeName)
      }
    }
    for f in messageProto.extension_p {
      if f.type == .group {
        protoNameIsGroup.insert(f.typeName)
      }
    }
    for e in messageProto.enumType {
      populateFrom(enumProto: e, prefix: name)
    }
    for m in messageProto.nestedType {
      populateFrom(messageProto: m, prefix: name)
    }
  }

  func getSwiftNameForEnumCase(path: String, caseName: String) -> String {
    return request.getSwiftNameForEnumCase(path: path, caseName: caseName)
  }

  func generateResponse() -> CodeGeneratorResponse {
    var response = CodeGeneratorResponse()
    let explicit = Set<String>(request.fileToGenerate)

    for fileProto in request.protoFile where explicit.contains(fileProto.name) {
      var printer = CodePrinter()
      let file = FileGenerator(descriptor: fileProto, generatorOptions: options)
      file.generateOutputFile(printer: &printer, context: self)
      let fileResponse = CodeGeneratorResponse.File(name: file.outputFilename, content: printer.content)
      response.file.append(fileResponse)
    }
    return response
  }
}