Add type annotations to entire function (#1138)

src/FsAutoComplete/CodeFixes/AddExplicitTypeAnnotation.fs

@@ -1,90 +1,156 @@
 module FsAutoComplete.CodeFix.AddExplicitTypeAnnotation
 
-open System
 open FsToolkit.ErrorHandling
 open FsAutoComplete.CodeFix.Types
 open Ionide.LanguageServerProtocol.Types
 open FsAutoComplete
 open FsAutoComplete.LspHelpers
-open FSharp.Compiler.CodeAnalysis
 open FSharp.Compiler.Symbols
 open FSharp.Compiler.Syntax
 open FSharp.Compiler.Text.Range
 open FsAutoComplete.Core.InlayHints
 open FsAutoComplete.Core
 
-
 let toLspEdit ({ Pos = insertAt; Text = text }: HintInsertion) =
   { Range = fcsPosToProtocolRange insertAt
     NewText = text }
 
 let toLspEdits (edits: HintInsertion[]) = edits |> Array.map toLspEdit
 
-[<Obsolete>] //TODO: correct?
-let private isPositionContainedInUntypedImplicitCtorParameter input pos =
-  let result =
+let title = "Add explicit type annotation"
+
+let rec nonTypedParameterName p =
+  match p with
+  | SynPat.Named(ident = SynIdent(ident, _)) -> Some ident
+  | SynPat.Paren(pat = p) -> nonTypedParameterName p
+  | _ -> None
+
+/// Captures a SynBinding that either has no return type or has parameters that are not typed.
+let (|FunctionBindingWithMissingTypes|_|) =
+  function
+  | SynBinding(
+      headPat = SynPat.LongIdent(longDotId = lid; argPats = SynArgPats.Pats parameters) as headPat
+      returnInfo = None
+      trivia = { LeadingKeyword = lk }) ->
+    let bindingStartRange = unionRanges lk.Range lid.Range
+    Some(bindingStartRange, Some headPat.Range, parameters.Length, List.choose nonTypedParameterName parameters)
+  | SynBinding(
+      headPat = SynPat.LongIdent(longDotId = lid; argPats = SynArgPats.Pats parameters)
+      returnInfo = Some _
+      trivia = { LeadingKeyword = lk }) ->
+    let bindingStartRange = unionRanges lk.Range lid.Range
+    let nonTypedParameters = List.choose nonTypedParameterName parameters
+
+    if List.isEmpty nonTypedParameters then
+      None
+    else
+      Some(bindingStartRange, None, parameters.Length, nonTypedParameters)
+  | _ -> None
+
+/// <summary>
+/// Try and find a SynBinding function where either the return type or any parameter is missing a type definition.
+/// </summary>
+/// <param name="parseAndCheck"></param>
+/// <param name="textDocument"></param>
+/// <param name="sourceText"></param>
+/// <param name="lineStr"></param>
+/// <param name="cursorPos">Expected to be between the start of the leading keyword and the end of the function name.</param>
+let tryFunctionIdentifier (parseAndCheck: ParseAndCheckResults) textDocument sourceText lineStr cursorPos =
+  let bindingInfo =
     SyntaxTraversal.Traverse(
-      pos,
-      input,
+      cursorPos,
+      parseAndCheck.GetAST,
       { new SyntaxVisitorBase<_>() with
-          member _.VisitModuleDecl(_, defaultTraverse, decl) =
-            match decl with
-            | SynModuleDecl.Types(typeDefns = typeDefns) ->
-              option {
-                let! ctorArgs =
-                  typeDefns
-                  |> List.tryPick (function
-                    | SynTypeDefn(implicitConstructor = Some(SynMemberDefn.ImplicitCtor(ctorArgs = args))) when
-                      rangeContainsPos args.Range pos
-                      ->
-                      Some args
-                    | _ -> None)
-
-                match ctorArgs with
-                | SynSimplePats.SimplePats(pats = pats) ->
-                  let! pat = pats |> List.tryFind (fun pat -> rangeContainsPos pat.Range pos)
-
-                  let rec tryGetUntypedIdent =
-                    function
-                    | SynSimplePat.Id(ident = ident) when rangeContainsPos ident.idRange pos -> Some ident
-                    | SynSimplePat.Attrib(pat = pat) when rangeContainsPos pat.Range pos -> tryGetUntypedIdent pat
-                    | SynSimplePat.Typed _
-                    | _ -> None
-
-                  return! tryGetUntypedIdent pat
-                | _ -> return! None
-              }
-              |> Option.orElseWith (fun _ -> defaultTraverse decl)
-            | _ -> defaultTraverse decl }
+          member _.VisitExpr(path, traverseSynExpr, defaultTraverse, expr) = defaultTraverse expr
+
+          member _.VisitBinding(path, defaultTraverse, binding) =
+            match binding with
+            | FunctionBindingWithMissingTypes(bindingStartRange,
+                                              headPatRangeOpt,
+                                              totalParameterCount,
+                                              nonTypedParameters) when rangeContainsPos bindingStartRange cursorPos ->
+              Some(bindingStartRange, headPatRangeOpt, totalParameterCount, nonTypedParameters)
+            | _ -> defaultTraverse binding }
     )
 
-  result.IsSome
-
-[<Obsolete>] //TODO: correct
-let private isSymbolToTriggerTypeAnnotation
-  (funcOrValue: FSharpMemberOrFunctionOrValue)
-  (symbolUse: FSharpSymbolUse)
-  (parseFileResults: FSharpParseFileResults)
-  =
-  (funcOrValue.IsValue
-   || (funcOrValue.IsFunction
-       && parseFileResults.IsBindingALambdaAtPosition symbolUse.Range.Start))
-  //TODO: check here for curried parameter? necessary? Or handled by `tryGetExplicitTypeInfo`?
-  && not funcOrValue.IsMember
-  && not funcOrValue.IsMemberThisValue
-  && not funcOrValue.IsConstructorThisValue
-  && not (PrettyNaming.IsOperatorDisplayName funcOrValue.DisplayName)
-
-
-let title = "Add explicit type annotation"
+  match bindingInfo with
+  | None -> []
+  | Some(bindingStartRange, headPatRangeOpt, untypedParameterCount, parameters) ->
+    match parseAndCheck.TryGetSymbolUse bindingStartRange.End lineStr with
+    | Some symbolUse ->
+      match symbolUse.Symbol with
+      | :? FSharpMemberOrFunctionOrValue as mfv when isPotentialTargetForTypeAnnotation true (symbolUse, mfv) ->
+        let returnTypeEdits =
+          match headPatRangeOpt with
+          | None -> [] // The return type is already present
+          | Some headPatRange ->
+            let returnTypeText =
+              if not mfv.FullType.IsFunctionType then
+                mfv.ReturnParameter.Type.Format(symbolUse.DisplayContext)
+              else
+                // We can't really be trust mfv.ReturnParameter, it will only contain the last type in a function type.
+                // Instead we collect all types and skip the amount of parameters we have in the function definition.
+                let allTypesFromFunctionType: FSharpType list =
+                  let rec visit (t: FSharpType) (continuation: FSharpType list -> FSharpType list) =
+                    if not t.IsFunctionType then
+                      continuation [ t ]
+                    else
+                      let funcType = t.GenericArguments.[0]
+                      let argType = t.GenericArguments.[1]
+
+                      if not argType.IsFunctionType then
+                        continuation [ funcType; argType ]
+                      else
+                        visit argType (fun types -> funcType :: types |> continuation)
+
+                  visit mfv.FullType id
+
+                if allTypesFromFunctionType.Length <= untypedParameterCount then
+                  mfv.ReturnParameter.Type.Format(symbolUse.DisplayContext)
+                else
+                  allTypesFromFunctionType
+                  |> List.skip untypedParameterCount
+                  |> List.map (fun t ->
+                    let formattedType = t.Format(symbolUse.DisplayContext)
+
+                    if t.IsFunctionType then
+                      $"({formattedType})"
+                    else
+                      formattedType)
+                  |> String.concat " -> "
+
+            // Put the return type after the current headPat.
+            [ { Range = fcsPosToProtocolRange headPatRange.End
+                NewText = $" : {returnTypeText}" } ]
+
+        let parameterEdits =
+          parameters
+          |> List.choose (fun ident ->
+            InlayHints.tryGetDetailedExplicitTypeInfo
+              (InlayHints.isPotentialTargetForTypeAnnotation true)
+              (sourceText, parseAndCheck)
+              ident.idRange.Start
+            |> Option.bind (fun (symbolUse, mfv, explTy) ->
+              explTy.TryGetTypeAndEdits(mfv.FullType, symbolUse.DisplayContext)
+              |> Option.map (fun (_, edits) -> toLspEdits edits)))
+          |> Seq.collect id
+          |> Seq.toArray
+
+        [ { File = textDocument
+            Title = title
+            Edits = [| yield! parameterEdits; yield! returnTypeEdits |]
+            Kind = FixKind.Refactor
+            SourceDiagnostic = None } ]
+      | _ -> []
+    | _ -> []
 
 let fix (getParseResultsForFile: GetParseResultsForFile) : CodeFix =
   fun codeActionParams ->
     asyncResult {
       let filePath = codeActionParams.TextDocument.GetFilePath() |> Utils.normalizePath
 
       let fcsStartPos = protocolPosToPos codeActionParams.Range.Start
-      let! (parseAndCheck, lineStr, sourceText) = getParseResultsForFile filePath fcsStartPos
+      let! parseAndCheck, lineStr, sourceText = getParseResultsForFile filePath fcsStartPos
 
       let res =
         InlayHints.tryGetDetailedExplicitTypeInfo
@@ -93,7 +159,14 @@ let fix (getParseResultsForFile: GetParseResultsForFile) : CodeFix =
           fcsStartPos
 
       match res with
-      | None -> return []
+      | None ->
+        return
+          tryFunctionIdentifier
+            parseAndCheck
+            codeActionParams.TextDocument
+            sourceText
+            lineStr
+            (protocolPosToPos codeActionParams.Range.End)
       | Some(symbolUse, mfv, explTy) ->
         match explTy.TryGetTypeAndEdits(mfv.FullType, symbolUse.DisplayContext) with
         | None -> return []

test/FsAutoComplete.Tests.Lsp/CodeFixTests/AddExplicitTypeAnnotationTests.fs

@@ -510,5 +510,120 @@ let tests state =
           """
           (Diagnostics.acceptAll)
           selectCodeFix
+
+      testCaseAsync "add type annotations to entire function" <|
+        CodeFix.check server
+          """
+          let x$0 y z =
+              ignore<int> y
+              ignore<char> z
+              0
+          """
+          Diagnostics.acceptAll
+          selectCodeFix
+          """
+          let x (y: int) (z: char) : int =
+              ignore<int> y
+              ignore<char> z
+              0
+          """
+
+      testCaseAsync "add type annotations to function with lambda function body" <|
+        CodeFix.check server
+          """
+          let f1$0 a = fun b -> a + b
+          """
+          Diagnostics.acceptAll
+          selectCodeFix
+          """
+          let f1 (a: int) : int -> int = fun b -> a + b
+          """
+
+      testCaseAsync "add type annotations to function with sequential lambda function body" <|
+        CodeFix.check server
+          """
+          let f2$0 a =
+              ()
+              fun b -> a + b
+          """
+          Diagnostics.acceptAll
+          selectCodeFix
+          """
+          let f2 (a: int) : int -> int =
+              ()
+              fun b -> a + b
+          """
+
+      testCaseAsync "add return type annotation when other parameters are typed" <|
+        CodeFix.check server
+          """
+          let f1$0 (a:int) = fun (b:char) -> System.TimeSpan.Zero
+          """
+          Diagnostics.acceptAll
+          selectCodeFix
+          """
+          let f1 (a:int) : char -> System.TimeSpan = fun (b:char) -> System.TimeSpan.Zero
+          """
+
+      testCaseAsync "add return type annotation for match lambda" <|
+        CodeFix.check server
+          """
+          let f1$0 = function | None -> 1 | Some _ -> 2
+          """
+          Diagnostics.acceptAll
+          selectCodeFix
+          """
+          let f1: 'a option -> int = function | None -> 1 | Some _ -> 2
+          """
+
+      testCaseAsync "add return type annotation when cursor is on let keyword" <|
+        CodeFix.check server
+          """
+          let$0 f g h = ignore<int> g ; ignore<string> h ; - 9.0
+          """
+          Diagnostics.acceptAll
+          selectCodeFix
+          """
+          let f (g: int) (h: string) : float = ignore<int> g ; ignore<string> h ; - 9.0
+          """
+
+      testCaseAsync "add type annotation for parameter when cursor is on function name" <|
+        CodeFix.check server
+          """
+          let f$0 (a:int) b : int = a + b
+          """
+          Diagnostics.acceptAll
+          selectCodeFix
+          """
+          let f (a:int) (b: int) : int = a + b
+          """
+
+      testCaseAsync "add type annotation for local function" <|
+        CodeFix.check server
+          """
+          do
+              let f$0 a b = a + b
+              ()
+          """
+          Diagnostics.acceptAll
+          selectCodeFix
+          """
+          do
+              let f (a: int) (b: int) : int = a + b
+              ()
+          """
+
+      testCaseAsync "add type annotation for recursive function" <|
+        CodeFix.check server
+          """
+          let rec a b = b - 1
+          and c$0 d e = d - e
+          """
+          Diagnostics.acceptAll
+          selectCodeFix
+          """
+          let rec a b = b - 1
+          and c (d: int) (e: int) : int = d - e
+          """
     ]
   ])