Merge pull request #71 from pontem-network/public-friend

don't use text to determine visibility modifier for the function
pontem-network · Sep 7, 2022 · 3de2dca · 3de2dca
2 parents 4c127cb + c53f770
commit 3de2dca
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 4 deletions.
diff --git a/src/main/kotlin/org/move/lang/core/psi/ext/MvFunction.kt b/src/main/kotlin/org/move/lang/core/psi/ext/MvFunction.kt
@@ -26,10 +26,10 @@ enum class FunctionVisibility {
 val MvFunction.visibility: FunctionVisibility
     get() {
         val visibility = this.functionVisibilityModifier ?: return FunctionVisibility.PRIVATE
-        return when (visibility.node.text) {
-            "public" -> FunctionVisibility.PUBLIC
-            "public(friend)" -> FunctionVisibility.PUBLIC_FRIEND
-            "public(script)" -> FunctionVisibility.PUBLIC_SCRIPT
+        return when {
+            visibility.hasChild(MvElementTypes.FRIEND) -> FunctionVisibility.PUBLIC_FRIEND
+            visibility.hasChild(MvElementTypes.SCRIPT_KW) -> FunctionVisibility.PUBLIC_SCRIPT
+            visibility.hasChild(MvElementTypes.PUBLIC) -> FunctionVisibility.PUBLIC
             else -> FunctionVisibility.PRIVATE
         }
     }

diff --git a/src/test/kotlin/org/move/lang/resolve/ResolveFunctionTest.kt b/src/test/kotlin/org/move/lang/resolve/ResolveFunctionTest.kt
@@ -469,4 +469,20 @@ class ResolveFunctionTest : ResolveTestCase() {
         use 0x1::caller::call;
     }
     """)
+
+    fun `test friend function from another module with spaces after public`() = checkByCode("""
+    module 0x1::A {
+        friend 0x1::B; 
+        public ( friend) fun call_a() {}
+                            //X
+    }        
+    module 0x1::B {
+        use 0x1::A;
+        
+        fun main() {
+            A::call_a();
+                  //^
+        }
+    }
+    """)
 }
diff --git a/src/test/resources/org/move/lang/parser/complete/friend.move b/src/test/resources/org/move/lang/parser/complete/friend.move
@@ -5,4 +5,5 @@ module M1 {
     friend 0x2::M2;
 
     public(friend) fun f_of_m1() {}
+    public ( friend ) fun f_of_m2() {}
 }
diff --git a/src/test/resources/org/move/lang/parser/complete/friend.txt b/src/test/resources/org/move/lang/parser/complete/friend.txt
@@ -51,5 +51,26 @@ FILE
         MvCodeBlockImpl(CODE_BLOCK)
           PsiElement({)('{')
           PsiElement(})('}')
+      PsiWhiteSpace('\n    ')
+      MvFunctionImpl(FUNCTION)
+        MvFunctionVisibilityModifierImpl(FUNCTION_VISIBILITY_MODIFIER)
+          PsiElement(public)('public')
+          PsiWhiteSpace(' ')
+          PsiElement(()('(')
+          PsiWhiteSpace(' ')
+          PsiElement(friend_kw)('friend')
+          PsiWhiteSpace(' ')
+          PsiElement())(')')
+        PsiWhiteSpace(' ')
+        PsiElement(fun)('fun')
+        PsiWhiteSpace(' ')
+        PsiElement(IDENTIFIER)('f_of_m2')
+        MvFunctionParameterListImpl(FUNCTION_PARAMETER_LIST)
+          PsiElement(()('(')
+          PsiElement())(')')
+        PsiWhiteSpace(' ')
+        MvCodeBlockImpl(CODE_BLOCK)
+          PsiElement({)('{')
+          PsiElement(})('}')
       PsiWhiteSpace('\n')
       PsiElement(})('}')