lambda type hints

pontem-network · Oct 24, 2024 · e68459d · e68459d
1 parent e1db952
commit e68459d
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 11 deletions.
diff --git a/src/main/kotlin/org/move/ide/hints/type/MvTypeInlayHintsProvider2.kt b/src/main/kotlin/org/move/ide/hints/type/MvTypeInlayHintsProvider2.kt
@@ -6,8 +6,10 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import org.move.ide.presentation.hintText
+import org.move.lang.core.psi.MvFunctionParameter
 import org.move.lang.core.psi.MvLetStmt
 import org.move.lang.core.psi.MvPatBinding
+import org.move.lang.core.psi.MvSchemaFieldStmt
 import org.move.lang.core.psi.ext.bindingOwner
 import org.move.lang.core.psi.ext.endOffset
 import org.move.lang.core.psi.ext.hasAncestor
@@ -31,8 +33,9 @@ class MvTypeInlayHintsProvider2: InlayHintsProvider {
             // skip private variables
             if (patBinding.name.startsWith("_")) return
 
-            // only show bindings for let statements
-            if (patBinding.bindingOwner !is MvLetStmt) return
+            // does not show hints for bindings with explicit type annotations
+            val owner = patBinding.bindingOwner
+            if (owner is MvFunctionParameter || owner is MvSchemaFieldStmt) return
 
             val contextInferenceOwner = patBinding.inferenceOwner() ?: return
 

diff --git a/src/main/kotlin/org/move/lang/core/psi/ext/MvBindingPat.kt b/src/main/kotlin/org/move/lang/core/psi/ext/MvBindingPat.kt
@@ -18,6 +18,7 @@ val MvPatBinding.bindingOwner: PsiElement?
         it is MvLetStmt
                 || it is MvFunctionParameter
                 || it is MvSchemaFieldStmt
+                || it is MvLambdaParameter
     }
 
 sealed class RsBindingModeKind {

diff --git a/src/main/kotlin/org/move/lang/core/types/infer/TypeInferenceWalker.kt b/src/main/kotlin/org/move/lang/core/types/infer/TypeInferenceWalker.kt
@@ -235,7 +235,7 @@ class TypeInferenceWalker(
             is MvDerefExpr -> inferDerefExprTy(expr)
             is MvLitExpr -> inferLitExprTy(expr, expected)
             is MvTupleLitExpr -> inferTupleLitExprTy(expr, expected)
-            is MvLambdaExpr -> inferLambdaExpr(expr, expected)
+            is MvLambdaExpr -> inferLambdaExprTy(expr, expected)
 
             is MvMoveExpr -> expr.expr?.inferType() ?: TyUnknown
             is MvCopyExpr -> expr.expr?.inferType() ?: TyUnknown
@@ -379,7 +379,7 @@ class TypeInferenceWalker(
         return TyReference(innerRefTy, mutability, ctx.msl)
     }
 
-    private fun inferLambdaExpr(lambdaExpr: MvLambdaExpr, expected: Expectation): Ty {
+    private fun inferLambdaExprTy(lambdaExpr: MvLambdaExpr, expected: Expectation): Ty {
         val bindings = lambdaExpr.parametersAsBindings
         val lambdaTy =
             (expected.onlyHasTy(this.ctx) as? TyLambda) ?: TyLambda.unknown(bindings.size)
@@ -398,10 +398,7 @@ class TypeInferenceWalker(
         val baseFuncTy =
             when (namedItem) {
                 is MvFunctionLike -> {
-                    val itemTy = instantiatePath<TyFunction>(path, namedItem) ?: return TyUnknown
-//                    val (itemTy, _) = instantiateMethodOrPath<TyFunction>(path, namedItem)
-//                        ?: return TyUnknown
-                    itemTy
+                    instantiatePath<TyFunction>(path, namedItem) ?: return TyUnknown
                 }
                 is MvFieldsOwner -> {
                     val tupleFields = namedItem.tupleFields
@@ -658,9 +655,9 @@ class TypeInferenceWalker(
         formalArgs: List<Ty>,
     ): List<Ty> {
         val resolvedFormalRet = resolveTypeVarsIfPossible(formalRet)
-        val retTy = expectedRet.onlyHasTy(ctx) ?: return emptyList()
+        val expectedRetTy = expectedRet.onlyHasTy(ctx) ?: return emptyList()
         return ctx.freezeUnification {
-            if (ctx.combineTypes(retTy, resolvedFormalRet).isOk) {
+            if (ctx.combineTypes(expectedRetTy, resolvedFormalRet).isOk) {
                 formalArgs.map { ctx.resolveTypeVarsIfPossible(it) }
             } else {
                 emptyList()

diff --git a/src/test/kotlin/org/move/ide/hints/InlayTypeHintsProvider2Test.kt b/src/test/kotlin/org/move/ide/hints/InlayTypeHintsProvider2Test.kt
@@ -117,7 +117,16 @@ class InlayTypeHintsProvider2Test: DeclarativeInlayHintsProviderTestCase() {
             }
         }
     """)
-
+
+    fun `test lambda expr`() = checkByText("""
+        module 0x1::m {
+            fun callback(elem: u8, ident: |u8|u8): u8 { ident(elem) }
+            fun main() {
+                callback(10, |elem/*<# : |u8 #>*/| elem + 1);
+            }
+        }        
+    """)
+
     private fun checkByText(@Language("Move") code: String) {
         doTestProvider("main.move", code, MvTypeInlayHintsProvider2())
     }

diff --git a/src/test/kotlin/org/move/lang/types/ExpressionTypesTest.kt b/src/test/kotlin/org/move/lang/types/ExpressionTypesTest.kt
@@ -2165,4 +2165,24 @@ module 0x1::main {
             }
         }        
     """)
+
+    fun `test lambda expr binding`() = testBinding("""
+        module 0x1::m {
+            fun callback(elem: u8, ident: |u8|u8): u8 { ident(elem) }
+            fun main() {
+                callback(10, |elem| elem + 1);
+                             //^ u8
+            }
+        }        
+    """)
+
+    fun `test lambda expr binding with inference`() = testBinding("""
+        module 0x1::m {
+            fun callback<E>(elem: E, ident: |E|E): E { ident(elem) }
+            fun main() {
+                callback(10u8, |elem| elem + 1);
+                               //^ u8
+            }
+        }        
+    """)
 }