{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE QuasiQuotes #-}

module Futhark.LSP.CodeAction (getCodeActions) where

import Data.Function ((&))
import Data.Loc (Pos (Pos), locOf)
import Data.Map qualified as M
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Set qualified as S
import Data.Text (Text)
import Data.Text qualified as T
import Futhark.LSP.State (State, getStaleMapping)
import Futhark.LSP.Tool (bindingsInState, filterByLspRange)
import Futhark.LSP.TypeAscription (TypeAscription (TypeAscLet, TypeAscParam, TypeAscReturn, TypeAscType), missingAscriptions, missingTypeParameters)
import Futhark.Util.Pretty (prettyText)
import Language.Futhark.Core (VName, baseText)
import Language.Futhark.FreeVars (freeInType, fvVars)
import Language.Futhark.Prop (typeVars)
import Language.Futhark.Syntax (RetTypeBase (RetType), TypeParamBase)
import Language.LSP.Protocol.Types (CodeAction (..), CodeActionKind (CodeActionKind_Custom), Command, Position (Position), Range (Range), TextEdit (..), Uri, WorkspaceEdit (..), type (|?) (InR))
import NeatInterpolation qualified as NI

getCodeActions :: Uri -> Range -> State -> FilePath -> [Command |? CodeAction]
getCodeActions :: Uri -> Range -> State -> FilePath -> [Command |? CodeAction]
getCodeActions Uri
file_uri Range
range State
state FilePath
filepath = [Command |? CodeAction]
-> Maybe [Command |? CodeAction] -> [Command |? CodeAction]
forall a. a -> Maybe a -> a
fromMaybe [] (Maybe [Command |? CodeAction] -> [Command |? CodeAction])
-> Maybe [Command |? CodeAction] -> [Command |? CodeAction]
forall a b. (a -> b) -> a -> b
$ do
  allBindings <- State -> Maybe (Map VName BoundTo)
bindingsInState State
state
  let inferredTypes = (BoundTo -> Map VName (Maybe Pos, TypeParamBase VName))
-> Map VName BoundTo -> Map VName (Maybe Pos, TypeParamBase VName)
forall m a. Monoid m => (a -> m) -> Map VName a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap BoundTo -> Map VName (Maybe Pos, TypeParamBase VName)
missingTypeParameters Map VName BoundTo
allBindings
  let mapping = State -> FilePath -> Maybe PositionMapping
getStaleMapping State
state FilePath
filepath
  bindings <-
    filterByLspRange mapping range filepath (locOf . snd) $
      M.assocs allBindings

  pure $
    concatMap
      ( mapMaybe (uncurry $ codeActions inferredTypes)
          . (\(VName
n, BoundTo
b) -> (VName
n,) (TypeAscription -> (VName, TypeAscription))
-> [TypeAscription] -> [(VName, TypeAscription)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BoundTo -> [TypeAscription]
missingAscriptions BoundTo
b)
      )
      bindings
      & map InR
  where
    codeActions ::
      M.Map VName (Maybe Pos, TypeParamBase VName) ->
      VName ->
      TypeAscription ->
      Maybe CodeAction
    codeActions :: Map VName (Maybe Pos, TypeParamBase VName)
-> VName -> TypeAscription -> Maybe CodeAction
codeActions Map VName (Maybe Pos, TypeParamBase VName)
types VName
name (TypeAscLet StructType
typ (Pos FilePath
_ Int
line Int
column Int
_)) =
      CodeAction -> Maybe CodeAction
forall a. a -> Maybe a
Just
        (CodeAction -> Maybe CodeAction) -> CodeAction -> Maybe CodeAction
forall a b. (a -> b) -> a -> b
$ Text -> [(Text, Int, Int)] -> CodeAction
forall {b}. Integral b => Text -> [(Text, b, b)] -> CodeAction
mkAction
          [NI.text|Insert type: `$nameText$typeText`|]
        ([(Text, Int, Int)] -> CodeAction)
-> [(Text, Int, Int)] -> CodeAction
forall a b. (a -> b) -> a -> b
$ (Text
typeText, Int
line, Int
column) (Text, Int, Int) -> [(Text, Int, Int)] -> [(Text, Int, Int)]
forall a. a -> [a] -> [a]
: Map VName (Maybe Pos, TypeParamBase VName)
-> Set VName -> [(Text, Int, Int)]
typeVarEdits Map VName (Maybe Pos, TypeParamBase VName)
types (StructType -> Set VName
forall {o}. TypeBase Size o -> Set VName
typesAndVars StructType
typ)
      where
        typeText :: Text
typeText = Text
": " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> StructType -> Text
forall a. Pretty a => a -> Text
prettyText StructType
typ
        nameText :: Text
nameText = VName -> Text
baseText VName
name
    codeActions Map VName (Maybe Pos, TypeParamBase VName)
types VName
name (TypeAscParam Pos
openPos StructType
typ Pos
pos) =
      CodeAction -> Maybe CodeAction
forall a. a -> Maybe a
Just
        (CodeAction -> Maybe CodeAction) -> CodeAction -> Maybe CodeAction
forall a b. (a -> b) -> a -> b
$ Text -> [(Text, Int, Int)] -> CodeAction
forall {b}. Integral b => Text -> [(Text, b, b)] -> CodeAction
mkAction
          [NI.text|Insert type: `($nameText: $typeText)`|]
        ([(Text, Int, Int)] -> CodeAction)
-> [(Text, Int, Int)] -> CodeAction
forall a b. (a -> b) -> a -> b
$ Pos -> Text -> (Text, Int, Int)
forall {a}. Pos -> a -> (a, Int, Int)
posEdit Pos
openPos Text
"("
          (Text, Int, Int) -> [(Text, Int, Int)] -> [(Text, Int, Int)]
forall a. a -> [a] -> [a]
: Pos -> Text -> (Text, Int, Int)
forall {a}. Pos -> a -> (a, Int, Int)
posEdit Pos
pos (Text
": " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
typeText Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")")
          (Text, Int, Int) -> [(Text, Int, Int)] -> [(Text, Int, Int)]
forall a. a -> [a] -> [a]
: Map VName (Maybe Pos, TypeParamBase VName)
-> Set VName -> [(Text, Int, Int)]
typeVarEdits Map VName (Maybe Pos, TypeParamBase VName)
types (StructType -> Set VName
forall {o}. TypeBase Size o -> Set VName
typesAndVars StructType
typ)
      where
        posEdit :: Pos -> a -> (a, Int, Int)
posEdit (Pos FilePath
_ Int
line Int
column Int
_) a
t = (a
t, Int
line, Int
column)
        typeText :: Text
typeText = StructType -> Text
forall a. Pretty a => a -> Text
prettyText StructType
typ
        nameText :: Text
nameText = VName -> Text
baseText VName
name
    codeActions Map VName (Maybe Pos, TypeParamBase VName)
types VName
name (TypeAscReturn ResRetType
typ (Pos FilePath
_ Int
line Int
column Int
_)) =
      CodeAction -> Maybe CodeAction
forall a. a -> Maybe a
Just
        (CodeAction -> Maybe CodeAction) -> CodeAction -> Maybe CodeAction
forall a b. (a -> b) -> a -> b
$ Text -> [(Text, Int, Int)] -> CodeAction
forall {b}. Integral b => Text -> [(Text, b, b)] -> CodeAction
mkAction
          [NI.text|Insert return type: `$nameText$typeText`|]
        ([(Text, Int, Int)] -> CodeAction)
-> [(Text, Int, Int)] -> CodeAction
forall a b. (a -> b) -> a -> b
$ (Text
typeText, Int
line, Int
column) (Text, Int, Int) -> [(Text, Int, Int)] -> [(Text, Int, Int)]
forall a. a -> [a] -> [a]
: Map VName (Maybe Pos, TypeParamBase VName)
-> Set VName -> [(Text, Int, Int)]
typeVarEdits Map VName (Maybe Pos, TypeParamBase VName)
types Set VName
retTypeVars
      where
        typeText :: Text
typeText = Text
" : " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> ResRetType -> Text
forall a. Pretty a => a -> Text
prettyText ResRetType
typ
        nameText :: Text
nameText = VName -> Text
baseText VName
name
        retTypeVars :: Set VName
retTypeVars =
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` TypeBase Size Freshness -> Set VName
forall {o}. TypeBase Size o -> Set VName
typesAndVars TypeBase Size Freshness
innerType
          where
            RetType [VName]
dims TypeBase Size Freshness
innerType = ResRetType
typ
    codeActions Map VName (Maybe Pos, TypeParamBase VName)
_ VName
_ (TypeAscType TypeParamBase VName
_ Pos
_) = Maybe CodeAction
forall a. Maybe a
Nothing

    typesAndVars :: TypeBase Size o -> Set VName
typesAndVars TypeBase Size o
t = TypeBase Size o -> Set VName
forall dim o. TypeBase dim o -> Set VName
typeVars TypeBase Size o
t Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` FV -> Set VName
fvVars (TypeBase Size o -> FV
forall o. TypeBase Size o -> FV
freeInType TypeBase Size o
t)

    -- Text edit information for referenced types that are not yet
    -- syntactically present
    typeVarEdits ::
      M.Map VName (Maybe Pos, TypeParamBase VName) ->
      S.Set VName ->
      [(Text, Int, Int)]
    typeVarEdits :: Map VName (Maybe Pos, TypeParamBase VName)
-> Set VName -> [(Text, Int, Int)]
typeVarEdits Map VName (Maybe Pos, TypeParamBase VName)
availableTypes Set VName
requiredTypes =
      Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
requiredTypes
        [VName]
-> ([VName] -> [(Maybe Pos, TypeParamBase VName)])
-> [(Maybe Pos, TypeParamBase VName)]
forall a b. a -> (a -> b) -> b
& (VName -> Maybe (Maybe Pos, TypeParamBase VName))
-> [VName] -> [(Maybe Pos, TypeParamBase VName)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Map VName (Maybe Pos, TypeParamBase VName)
availableTypes Map VName (Maybe Pos, TypeParamBase VName)
-> VName -> Maybe (Maybe Pos, TypeParamBase VName)
forall k a. Ord k => Map k a -> k -> Maybe a
M.!?)
        [(Maybe Pos, TypeParamBase VName)]
-> ([(Maybe Pos, TypeParamBase VName)] -> [(Text, Int, Int)])
-> [(Text, Int, Int)]
forall a b. a -> (a -> b) -> b
& ((Maybe Pos, TypeParamBase VName) -> Maybe (Text, Int, Int))
-> [(Maybe Pos, TypeParamBase VName)] -> [(Text, Int, Int)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Maybe Pos, TypeParamBase VName) -> Maybe (Text, Int, Int)
typeTriple
      where
        typeTriple :: (Maybe Pos, TypeParamBase VName) -> Maybe (Text, Int, Int)
        typeTriple :: (Maybe Pos, TypeParamBase VName) -> Maybe (Text, Int, Int)
typeTriple (Maybe Pos
pos, TypeParamBase VName
typ) = do
          Pos _ line col _ <- Maybe Pos
pos
          pure (" " <> prettyText typ, line, col)

    mkAction :: Text -> [(Text, b, b)] -> CodeAction
mkAction Text
title [(Text, b, b)]
edits =
      CodeAction
        { _isPreferred :: Maybe Bool
_isPreferred = Maybe Bool
forall a. Maybe a
Nothing,
          _disabled :: Maybe CodeActionDisabled
_disabled = Maybe CodeActionDisabled
forall a. Maybe a
Nothing,
          _diagnostics :: Maybe [Diagnostic]
_diagnostics = Maybe [Diagnostic]
forall a. Maybe a
Nothing,
          _edit :: Maybe WorkspaceEdit
_edit =
            WorkspaceEdit -> Maybe WorkspaceEdit
forall a. a -> Maybe a
Just (WorkspaceEdit -> Maybe WorkspaceEdit)
-> WorkspaceEdit -> Maybe WorkspaceEdit
forall a b. (a -> b) -> a -> b
$
              WorkspaceEdit
                { _documentChanges :: Maybe
  [TextDocumentEdit |? (CreateFile |? (RenameFile |? DeleteFile))]
_documentChanges = Maybe
  [TextDocumentEdit |? (CreateFile |? (RenameFile |? DeleteFile))]
forall a. Maybe a
Nothing,
                  _changeAnnotations :: Maybe (Map ChangeAnnotationIdentifier ChangeAnnotation)
_changeAnnotations = Maybe (Map ChangeAnnotationIdentifier ChangeAnnotation)
forall a. Maybe a
Nothing,
                  _changes :: Maybe (Map Uri [TextEdit])
_changes =
                    Map Uri [TextEdit] -> Maybe (Map Uri [TextEdit])
forall a. a -> Maybe a
Just (Map Uri [TextEdit] -> Maybe (Map Uri [TextEdit]))
-> Map Uri [TextEdit] -> Maybe (Map Uri [TextEdit])
forall a b. (a -> b) -> a -> b
$
                      Uri -> [TextEdit] -> Map Uri [TextEdit]
forall k a. k -> a -> Map k a
M.singleton
                        Uri
file_uri
                        (((Text, b, b) -> TextEdit) -> [(Text, b, b)] -> [TextEdit]
forall a b. (a -> b) -> [a] -> [b]
map (Text, b, b) -> TextEdit
forall {b}. Integral b => (Text, b, b) -> TextEdit
mkEdit [(Text, b, b)]
edits)
                },
          _title :: Text
_title = Text
title,
          _command :: Maybe Command
_command = Maybe Command
forall a. Maybe a
Nothing,
          _data_ :: Maybe Value
_data_ = Maybe Value
forall a. Maybe a
Nothing,
          _kind :: Maybe CodeActionKind
_kind = CodeActionKind -> Maybe CodeActionKind
forall a. a -> Maybe a
Just (CodeActionKind -> Maybe CodeActionKind)
-> CodeActionKind -> Maybe CodeActionKind
forall a b. (a -> b) -> a -> b
$ Text -> CodeActionKind
CodeActionKind_Custom Text
"TypeAscription"
        }

    mkEdit :: (Text, b, b) -> TextEdit
mkEdit (Text
text, b
line, b
column) =
      TextEdit
        { _range :: Range
_range =
            let toLsp :: b -> UInt
toLsp = b -> UInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral (b -> UInt) -> (b -> b) -> b -> UInt
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> b
forall a. Enum a => a -> a
pred
                insertPos :: Position
insertPos = UInt -> UInt -> Position
Position (b -> UInt
toLsp b
line) (b -> UInt
toLsp b
column)
             in Position -> Position -> Range
Range Position
insertPos Position
insertPos,
          _newText :: Text
_newText = Text -> Text
normalizeNames Text
text
        }
      where
        -- replace subscript characters
        normalizeNames :: Text -> Text
normalizeNames = (Char -> Char) -> Text -> Text
T.map ((Char -> Char) -> Text -> Text) -> (Char -> Char) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ \case
          Char
'\8320' -> Char
'0'
          Char
'\8321' -> Char
'1'
          Char
'\8322' -> Char
'2'
          Char
'\8323' -> Char
'3'
          Char
'\8324' -> Char
'4'
          Char
'\8325' -> Char
'5'
          Char
'\8326' -> Char
'6'
          Char
'\8327' -> Char
'7'
          Char
'\8328' -> Char
'8'
          Char
'\8329' -> Char
'9'
          Char
x -> Char
x