module Futhark.LSP.TypeAscription
  ( TypeAscription (..),
    missingAscriptions,
    missingTypeParameters,
  )
where

import Control.Arrow ((&&&))
import Data.Foldable (asum)
import Data.Function ((&))
import Data.Loc (Loc (..), Pos, locOf)
import Data.Map qualified as M
import Language.Futhark.Query
  ( BoundTo (..),
    TermBindSrc (..),
    TermBinding (..),
    TermFunData (..),
  )
import Language.Futhark.Syntax
  ( ResRetType,
    StructType,
    TypeParamBase,
    VName,
    typeParamName,
  )

data TypeAscription
  = -- | type hint for a let-binding or nested parameter binding
    TypeAscLet StructType Pos
  | -- | type hint for a parameter
    TypeAscParam Pos StructType Pos
  | -- | type hint for a return type
    TypeAscReturn ResRetType Pos
  | -- | type hint for inferred type parameters
    TypeAscType (TypeParamBase VName) Pos

-- techincally, this name is not entirely correct since it covers type
-- parameters as well
missingAscriptions :: BoundTo -> [TypeAscription]
missingAscriptions :: BoundTo -> [TypeAscription]
missingAscriptions (BoundTerm TermBinding
term (Loc Pos
termStart Pos
termEnd)) =
  case TermBinding
term of
    TermBinding
TermSize -> []
    TermVar TermBindSrc
_ StructType
_ (Just TypeAscription
_) -> []
    TermVar TermBindSrc
src StructType
inferredType Maybe TypeAscription
Nothing ->
      let bareHints :: [TypeAscription]
bareHints = [StructType -> Pos -> TypeAscription
TypeAscLet StructType
inferredType Pos
termEnd]
       in case TermBindSrc
src of
            TermBindSrc
TermBindId -> []
            TermBindSrc
TermBindLet -> [TypeAscription]
bareHints
            TermBindSrc
TermBindPat ->
              [Pos -> StructType -> Pos -> TypeAscription
TypeAscParam Pos
termStart StructType
inferredType Pos
termEnd]
            TermBindSrc
TermBindNested -> [TypeAscription]
bareHints
    TermFun TermFunData
tfData ->
      -- ordering is relevant, see the lsp documentation for inlay hints
      let isSynthesized :: a -> Bool
isSynthesized a
typeParam = case a -> Loc
forall a. Located a => a -> Loc
locOf a
typeParam of
            Loc
NoLoc -> Bool
True
            Loc
_ -> Bool
False
          inferredTypeParams :: [TypeParamBase VName]
inferredTypeParams = (TypeParamBase VName -> Bool)
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. (a -> Bool) -> [a] -> [a]
filter TypeParamBase VName -> Bool
forall {a}. Located a => a -> Bool
isSynthesized (TermFunData -> [TypeParamBase VName]
termFunTypeParams TermFunData
tfData)
          paramInfo :: TypeParamBase VName -> [TypeAscription] -> [TypeAscription]
paramInfo TypeParamBase VName
p = case TermFunData -> Maybe Pos
termFunNameEnd TermFunData
tfData of
            Maybe Pos
Nothing -> [TypeAscription] -> [TypeAscription]
forall a. a -> a
id
            Just Pos
pos -> (TypeParamBase VName -> Pos -> TypeAscription
TypeAscType TypeParamBase VName
p Pos
pos TypeAscription -> [TypeAscription] -> [TypeAscription]
forall a. a -> [a] -> [a]
:)
       in (TypeParamBase VName
 -> ([TypeAscription] -> [TypeAscription])
 -> [TypeAscription]
 -> [TypeAscription])
-> ([TypeAscription] -> [TypeAscription])
-> [TypeParamBase VName]
-> [TypeAscription]
-> [TypeAscription]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\TypeParamBase VName
p [TypeAscription] -> [TypeAscription]
f [TypeAscription]
hs -> TypeParamBase VName -> [TypeAscription] -> [TypeAscription]
paramInfo TypeParamBase VName
p ([TypeAscription] -> [TypeAscription])
-> [TypeAscription] -> [TypeAscription]
forall a b. (a -> b) -> a -> b
$ [TypeAscription] -> [TypeAscription]
f [TypeAscription]
hs) [TypeAscription] -> [TypeAscription]
forall a. a -> a
id [TypeParamBase VName]
inferredTypeParams ([TypeAscription] -> [TypeAscription])
-> [TypeAscription] -> [TypeAscription]
forall a b. (a -> b) -> a -> b
$
            let retType :: ResRetType
retType = TermFunData -> ResRetType
termFunRetType TermFunData
tfData
                retTypePos :: Maybe Pos
retTypePos =
                  [Maybe Pos] -> Maybe Pos
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum ([Maybe Pos] -> Maybe Pos) -> [Maybe Pos] -> Maybe Pos
forall a b. (a -> b) -> a -> b
$
                    ((TermFunData -> Maybe Pos) -> Maybe Pos)
-> [TermFunData -> Maybe Pos] -> [Maybe Pos]
forall a b. (a -> b) -> [a] -> [b]
map
                      ((TermFunData -> Maybe Pos) -> TermFunData -> Maybe Pos
forall a b. (a -> b) -> a -> b
$ TermFunData
tfData)
                      [ TermFunData -> Maybe Pos
termFunArgEnd,
                        TermFunData -> Maybe Pos
termFunTypeArgEnd,
                        TermFunData -> Maybe Pos
termFunNameEnd
                      ]
             in case (TermFunData -> Maybe TypeAscription
termFunAscription TermFunData
tfData, Maybe Pos
retTypePos) of
                  (Just TypeAscription
_, Maybe Pos
_) -> []
                  (Maybe TypeAscription
Nothing, Just Pos
pos) -> [ResRetType -> Pos -> TypeAscription
TypeAscReturn ResRetType
retType Pos
pos]
                  (Maybe TypeAscription, Maybe Pos)
_ -> []
missingAscriptions BoundTo
_ = []

missingTypeParameters :: BoundTo -> M.Map VName (Maybe Pos, TypeParamBase VName)
missingTypeParameters :: BoundTo -> Map VName (Maybe Pos, TypeParamBase VName)
missingTypeParameters (BoundTerm (TermFun TermFunData
tfData) Loc
_) =
  [(VName, (Maybe Pos, TypeParamBase VName))]
-> Map VName (Maybe Pos, TypeParamBase VName)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(VName, (Maybe Pos, TypeParamBase VName))]
inferredTypeParams
  where
    inferredTypeParams :: [(VName, (Maybe Pos, TypeParamBase VName))]
inferredTypeParams =
      TermFunData -> [TypeParamBase VName]
termFunTypeParams TermFunData
tfData
        [TypeParamBase VName]
-> ([TypeParamBase VName] -> [TypeParamBase VName])
-> [TypeParamBase VName]
forall a b. a -> (a -> b) -> b
& (TypeParamBase VName -> Bool)
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. (a -> Bool) -> [a] -> [a]
filter TypeParamBase VName -> Bool
forall {a}. Located a => a -> Bool
keepNoLoc
        [TypeParamBase VName]
-> ([TypeParamBase VName]
    -> [(VName, (Maybe Pos, TypeParamBase VName))])
-> [(VName, (Maybe Pos, TypeParamBase VName))]
forall a b. a -> (a -> b) -> b
& (TypeParamBase VName -> (VName, (Maybe Pos, TypeParamBase VName)))
-> [TypeParamBase VName]
-> [(VName, (Maybe Pos, TypeParamBase VName))]
forall a b. (a -> b) -> [a] -> [b]
map (TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName (TypeParamBase VName -> VName)
-> (TypeParamBase VName -> (Maybe Pos, TypeParamBase VName))
-> TypeParamBase VName
-> (VName, (Maybe Pos, TypeParamBase VName))
forall b c c'. (b -> c) -> (b -> c') -> b -> (c, c')
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& (TermFunData -> Maybe Pos
termFunNameEnd TermFunData
tfData,))
    keepNoLoc :: a -> Bool
keepNoLoc a
located = case a -> Loc
forall a. Located a => a -> Loc
locOf a
located of
      Loc
NoLoc -> Bool
True
      Loc
_ -> Bool
False
missingTypeParameters BoundTo
_ = Map VName (Maybe Pos, TypeParamBase VName)
forall k a. Map k a
M.empty