如何创建类似@Compose的函数调用限制注解@WithResource

一、实现@WithResource的调用上下文检查 #

要实现这种调用上下文的强制限制，普通注解无法完成，必须借助Kotlin编译期插件（或KSP结合编译插件）来做静态代码检查，以下是具体步骤：

1. 定义@WithResource注解 #

首先创建注解，标记需要限制的函数和Lambda：

@Target(AnnotationTarget.FUNCTION, AnnotationTarget.LAMBDA)
@Retention(AnnotationRetention.BINARY) // 保留到二进制，供编译插件读取
annotation class WithResource

2. 实现Kotlin编译插件 #

核心逻辑是在编译阶段扫描所有函数调用，对被@WithResource标记的函数，检查其调用上下文是否合法：

• 调用点所在的函数/Lambda必须也被@WithResource标记；
• 或者调用点位于openResource函数的@WithResource类型Lambda参数内部。

插件的关键检查逻辑伪代码（基于Kotlin Compiler API）：

// 简化的插件检查逻辑
fun checkCallContext(call: CallExpression, annotationFinder: AnnotationFinder) {
    val calledFunction = call.resolveToFunction()
    if (calledFunction.hasAnnotation<WithResource>()) {
        // 向上回溯调用所在的上下文
        var currentContext = call.parent
        var isInValidContext = false
        while (currentContext != null) {
            when (currentContext) {
                is FunctionDeclaration -> {
                    isInValidContext = currentContext.hasAnnotation<WithResource>()
                    break
                }
                is LambdaExpression -> {
                    val lambdaType = currentContext.resolveToType()
                    // 检查Lambda是否是openResource的@WithResource参数
                    isInValidContext = lambdaType.annotations.any { it is WithResource }
                    // 同时确认这个Lambda属于openResource函数
                    val parentCall = currentContext.parent as? CallExpression
                    isInValidContext = isInValidContext && parentCall?.resolveToFunction()?.name == "openResource"
                    if (isInValidContext) break
                }
            }
            currentContext = currentContext.parent
        }
        if (!isInValidContext) {
            throw CompilationError("Function ${calledFunction.name} can only be called in @WithResource context or openResource block")
        }
    }
}

3. 合法/非法调用示例 #

合法调用

// 定义资源管理函数，其Lambda参数标记@WithResource
fun <T> openResource(block: @WithResource (Resource) -> T): T {
    val resource = Resource() // 假设Resource构造函数私有，仅此处可创建
    return try {
        block(resource)
    } finally {
        resource.close()
    }
}

@WithResource
fun useResourceDetail(res: Resource) {
    // 合法：自身在@WithResource函数内
    readResource(res)
}

@WithResource
fun readResource(res: Resource) { /* 资源读取逻辑 */ }

fun main() {
    openResource { res ->
        // 合法：在openResource的@WithResource Lambda内
        useResourceDetail(res)
        readResource(res)
    }
}

非法调用（编译报错）

fun main() {
    val res = Resource() // 假设构造函数私有，此处本就无法编译
    // 即使能拿到res，编译插件会报错：不在合法上下文
    readResource(res)
}

二、防止Resource泄露的思路 #

结合编译期检查和类型约束，从根源上避免泄露：

1. 类型层面限制Resource的可见性 #

• 将Resource的构造函数设为私有，仅允许openResource函数内部创建实例，外部无法直接获取；
• 可选：给Resource添加一个内部标记接口，比如internal interface RestrictedResource，让Resource实现它，同时openResource的Lambda参数类型设为RestrictedResource，外部无法持有该类型的引用（若用internal修饰，仅限模块内，结合编译插件可进一步限制）。

2. 编译期检查泄露行为 #

扩展上述的编译插件，添加额外检查规则：

• 禁止将Resource实例赋值给非局部变量（比如外部类的属性、全局变量）；
• 禁止被@WithResource标记的函数返回Resource实例；
• 禁止将Resource作为参数传递给未被@WithResource标记的函数。

3. 运行时兜底 #

给Resource实现AutoCloseable接口，配合openResource的try-finally确保资源被关闭，即使编译期检查出现遗漏，也能在运行时回收资源：

class Resource : AutoCloseable {
    private constructor() // 私有构造
    companion object {
        // 仅openResource可调用的创建方法
        internal fun create(): Resource = Resource()
    }

    override fun close() {
        // 资源释放逻辑
    }
}

// 修改openResource，使用内部创建方法
fun <T> openResource(block: @WithResource (Resource) -> T): T {
    val resource = Resource.create()
    return try {
        block(resource)
    } finally {
        resource.close()
    }
}

内容的提问来源于stack exchange，提问作者Arthur Rainbow