override fun consume()

in sample/experiments-compiler/experiment-enums-compiler/kotlin/src/main/kotlin/com/uber/crumb/sample/experimentenumscompiler/ExperimentsCompiler.kt [81:158]


  override fun consume(context: CrumbContext,
      type: TypeElement,
      annotations: Collection<AnnotationMirror>,
      metadata: Set<ConsumerMetadata>) {

    // Must be a type that supports extension values
    if (type.kind != CLASS) {
      context.processingEnv
          .messager
          .printMessage(ERROR,
              "@${ExperimentsCollector::class.java.simpleName} is only applicable on classes when consuming!",
              type)
      return
    }

    val kmetadata = type.kotlinMetadata

    if (kmetadata !is KotlinClassMetadata) {
      context.processingEnv
          .messager
          .printMessage(ERROR,
              "@${ExperimentsCollector::class.java.simpleName} can't be applied to $type: must be a class. KMetadata was $kmetadata and annotations were [${type.annotationMirrors.joinToString { it.annotationType.asElement().simpleName }}]",
              type)
      return
    }

    val classData = kmetadata.data
    val (nameResolver, classProto) = classData

    // Must be an abstract class because we're generating the backing implementation.
    if (classProto.classKind != ProtoBuf.Class.Kind.OBJECT) {
      context.processingEnv
          .messager
          .printMessage(ERROR,
              "@${ExperimentsCollector::class.java.simpleName} can't be applied to $type: must be a Kotlin object class",
              type)
      return
    }

    val packageName = nameResolver.getString(classProto.fqName)
        .substringBeforeLast('/')
        .replace('/', '.')

    // Map of enum TypeElement to its members
    val experimentClasses = metadata
        .mapNotNull { it[METADATA_KEY] }
        .map { enumClass -> context.processingEnv.elementUtils.getTypeElement(enumClass) }
        .associate {
          it to it.enclosedElements
              .filter { it.kind == ENUM_CONSTANT }
              .map(Element::toString)
        }

    val initializerCode = experimentClasses
        .map { "%T::class.java to listOf(${it.value.joinToString(", ") { "%S" }})" }
        .joinToString()
    val initializerValues = experimentClasses
        .flatMap { listOf(it.key.asClassName(), *it.value.toTypedArray()) }
        .toTypedArray()
    val mapFunction = FunSpec.builder("experiments")
        .receiver(type.asClassName())
        .returns(Map::class.asClassName().parameterizedBy(
            Class::class.asClassName().parameterizedBy(
                WildcardTypeName.producerOf(
                    Enum::class.asClassName().parameterizedBy(STAR))),
            List::class.asClassName().parameterizedBy(
                String::class.asTypeName())))
        .addStatement("return mapOf($initializerCode)", *initializerValues)
        .build()

    // Generate the file
    val generatedDir = context.processingEnv.options[kaptGeneratedOption]?.let(::File)
        ?: throw IllegalStateException("Could not resolve kotlin generated directory!")
    FileSpec.builder(packageName, "${type.simpleName}_Experiments")
        .addFunction(mapFunction)
        .build()
        .writeTo(generatedDir)
  }