arrow.meta.phases.analysis.KtUtils.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of arrow-meta Show documentation
Functional companion to Kotlin's Standard Library
There is a newer version: 1.6.3-alpha.2
package arrow.meta.phases.analysis

import arrow.meta.ArrowMetaConfigurationKeys
import org.jetbrains.kotlin.builtins.KotlinBuiltIns
import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
import org.jetbrains.kotlin.builtins.isBuiltinFunctionalType
import org.jetbrains.kotlin.com.intellij.psi.PsiElement
import org.jetbrains.kotlin.com.intellij.psi.SyntaxTraverser
import org.jetbrains.kotlin.config.CompilerConfiguration
import org.jetbrains.kotlin.descriptors.CallableDescriptor
import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
import org.jetbrains.kotlin.platform.CommonPlatforms
import org.jetbrains.kotlin.psi.KtAnnotated
import org.jetbrains.kotlin.psi.KtBlockExpression
import org.jetbrains.kotlin.psi.KtClass
import org.jetbrains.kotlin.psi.KtDeclarationWithBody
import org.jetbrains.kotlin.psi.KtElement
import org.jetbrains.kotlin.psi.KtExpression
import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.psi.KtObjectDeclaration
import org.jetbrains.kotlin.psi.KtReturnExpression
import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
import org.jetbrains.kotlin.psi.psiUtil.astReplace
import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
import org.jetbrains.kotlin.resolve.descriptorUtil.platform
import org.jetbrains.kotlin.resolve.multiplatform.isCommonSource
import org.jetbrains.kotlin.resolve.scopes.MemberScope
import org.jetbrains.kotlin.types.KotlinType
import org.jetbrains.kotlin.types.TypeConstructor
import org.jetbrains.kotlin.types.TypeProjection
import org.jetbrains.kotlin.utils.addToStdlib.safeAs

fun KtDeclarationWithBody.body(): KtExpression? = bodyExpression ?: bodyBlockExpression

fun KtExpression.bodySourceAsExpression(): String? =
  when (this) {
    is KtBlockExpression ->
      statements
        .map {
          when (it) {
            is KtReturnExpression -> it.returnedExpression?.text
            else -> text
          }
        }
        .joinToString("\n")
        .drop(1)
        .dropLast(1)
    else -> text
  }

fun KtElement.transform(f: (KtElement) -> KtElement?): KtElement {
  accept(
    object : KtTreeVisitorVoid() {
      override fun visitKtElement(element: KtElement) {
        val result = f(element)
        if (result != null) {
          element.astReplace(result)
        }
        super.visitKtElement(element)
      }
    }
  )
  return this
}

fun KtElement.dfs(f: (KtElement) -> Boolean): List {
  val found = arrayListOf()
  accept(
    object : KtTreeVisitorVoid() {
      override fun visitKtElement(element: KtElement) {
        val result = f(element)
        if (result) found.add(element)
        super.visitKtElement(element)
      }
    }
  )
  return found
}

/**
 * collects all distinct witnesses of [f] from the receiver, where the return type contains pairs of
 * [A] and a list of other corresponding elements that full fill f
 */
fun  List.exists(f: (A, A) -> Boolean): List>> =
  fold(emptyList()) { acc: List>>, a: A ->
    acc + (a to filter { b: A -> if (a != b) f(a, b) else false })
  }

/**
 * traverse and filters starting from the root node [this] down to all it's children and applying
 * [f]
 */
fun  PsiElement.traverseFilter(on: Class, f: (A) -> B?): List =
  SyntaxTraverser.psiTraverser(this).filter(on).mapNotNull(f).toList()

/**
 * a convenient function that collects all child nodes [A] starting from [this] it applies
 * [traverseFilter] with the identity function
 */
fun  PsiElement.sequence(on: Class): List = traverseFilter(on) { it }

interface Eq { // from arrow
  fun A.eqv(other: A): Boolean
  fun A.neqv(other: A): Boolean = !eqv(other)

  companion object {
    inline operator fun  invoke(crossinline feqv: (A, A) -> Boolean): Eq =
      object : Eq {
        override fun A.eqv(other: A): Boolean = feqv(this, other)
      }
  }
}

/** defines Equality on the type constructor */
fun typeConstructorEq(): Eq = Eq { t1, t2 -> t1.constructor == t2.constructor }

/** defines Equality on types, where FunctionTypes are reduced to their return type */
fun resolveFunctionTypeEq(): Eq = Eq { t1, t2 ->
  resolveFunctionType(t1) == resolveFunctionType(t2)
}

/**
 * Given [eq] this function returns [KotlinType]s that [intersect] with the returnType from the list
 * in [types]. One concrete example for equality on [TypeConstructor] may look like this:
 * ```kotlin:ank
 * import org.jetbrains.kotlin.descriptors.CallableDescriptor
 * import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 * import org.jetbrains.kotlin.types.KotlinType
 * import org.jetbrains.kotlin.resolve.descriptorUtil.builtIns
 * import arrow.meta.phases.analysis.Eq
 * import arrow.meta.phases.analysis.intersect
 * //sampleStart
 *
 * fun typeConstructorEq(): Eq =
 *  Eq { t1, t2 ->
 *   t1.constructor == t2.constructor
 *  }
 *
 * /**
 * * this function yields true if the type constructor - short [TC] - of the returnType in F is equal to one TC in [types]
 * */
 * fun  F.returns( //
 *   types: KotlinBuiltIns.() -> List
 * ): Boolean =
 *   intersect(typeConstructorEq(), types).isNotEmpty()
 * //sampleEnd
 * ```
 * @see [org.jetbrains.kotlin.types.TypeUtils] for more abstractions
 * @param eq can be define for e.g.: [TypeConstructor], [MemberScope] or typeArguments List<
 * [TypeProjection]>, etc.
 * @see functionTypeEq
 */
fun  C.intersect(
  eq: Eq,
  types: KotlinBuiltIns.() -> List
): List =
  eq.run {
    returnType?.let { result: KotlinType -> builtIns.types().filter { it.eqv(result) } }
      ?: emptyList()
  }

/**
 * given [eq] this function returns a List of [KotlinType] that are contained by both [list] and
 * [other]
 * @param eq can be defined for [TypeConstructor], [MemberScope] or typeArguments List<
 * [TypeProjection]>, etc.
 * @see intersect
 */
fun  D.intersect(
  eq: Eq,
  list: List,
  other: KotlinBuiltIns.() -> List
): List =
  eq.run {
    val set = list.toMutableList()
    set.retainAll { t1 -> builtIns.other().any { t2 -> t1.eqv(t2) } }
    set.toList()
  }

/** resolves FunctionType to it's returnType */
val resolveFunctionType: (KotlinType) -> KotlinType
  get() = { if (it.isBuiltinFunctionalType) it.getReturnTypeFromFunctionType() else it }

/** naive type equality where function types are reduced to their return type */
val returnTypeEq: Eq
  get() = Eq { a, b -> resolveFunctionType(a) == resolveFunctionType(b) }

fun KtAnnotated.isAnnotatedWith(regex: Regex): Boolean =
  annotationEntries.any { it.text.matches(regex) }

val KtClass.companionObject: KtObjectDeclaration?
  get() =
    declarations
      .singleOrNull { it.safeAs()?.isCompanion() == true }
      .safeAs()

fun DeclarationDescriptor.skipGeneration(): Boolean =
  platform != CommonPlatforms.defaultCommonPlatform &&
    (findPsi()?.containingFile as? KtFile)?.isCommonSource == true

fun getOrCreateBaseDirectory(configuration: CompilerConfiguration?): java.io.File {
  val parentBuildPath: String? =
    configuration?.get(ArrowMetaConfigurationKeys.GENERATED_SRC_OUTPUT_DIR)?.firstOrNull()
      ?: System.getProperty("arrow.meta.generate.source.dir")
  checkNotNull(parentBuildPath) {
    "Generated sources output dir is not found: ${configuration?.get(ArrowMetaConfigurationKeys.GENERATED_SRC_OUTPUT_DIR)} ${System.getProperty("arrow.meta.generate.source.dir")}"
  }
  val directory = java.io.File("$parentBuildPath")
  directory.mkdirs()
  return directory
}