Include LHS of binary exprs in annotation (#3640)

* Include LHS of binary exprs in annotation

* Add combinators from Hover branch for testing

* Finish adding LSP annotation tests

* Fix custom op lhs annotations

codebase2/core/U/Core/ABT.hs

@@ -17,7 +17,7 @@ data ABT f v r
   | Cycle r
   | Abs v r
   | Tm (f r)
-  deriving (Show, Functor, Foldable, Traversable)
+  deriving stock (Eq, Show, Functor, Foldable, Traversable)
 
 -- | At each level in the tree, we store the set of free variables and
 -- a value of type `a`. Variables are of type `v`.

lib/unison-prelude/src/Unison/Prelude.hs

@@ -18,6 +18,8 @@ module Unison.Prelude
     whenJustM,
     eitherToMaybe,
     maybeToEither,
+    altSum,
+    altMap,
 
     -- * @Either@ control flow
     onLeft,
@@ -74,6 +76,14 @@ import UnliftIO as X (MonadUnliftIO (..), askRunInIO, askUnliftIO, try, withUnli
 import qualified UnliftIO
 import Witherable as X (filterA, forMaybe, mapMaybe, wither, witherMap)
 
+-- | Like 'fold' but for Alternative.
+altSum :: (Alternative f, Foldable t) => t (f a) -> f a
+altSum = foldl' (<|>) empty
+
+-- | Like 'foldMap' but for Alternative.
+altMap :: (Alternative f, Foldable t) => (a -> f b) -> t a -> f b
+altMap f = altSum . fmap f . toList
+
 -- | E.g.
 --
 -- @@

parser-typechecker/src/Unison/Codebase/Init/CreateCodebaseError.hs

@@ -10,3 +10,4 @@ type Pretty = P.Pretty P.ColorText
 
 data CreateCodebaseError
   = CreateCodebaseAlreadyExists
+  deriving stock (Show)

parser-typechecker/src/Unison/Syntax/TermParser.hs

@@ -934,11 +934,11 @@ infixAppOrBooleanOp :: Var v => TermP v
 infixAppOrBooleanOp = chainl1 term4 (or <|> and <|> infixApp)
   where
     or = orf <$> label "or" (reserved "||")
-    orf op lhs rhs = Term.or (ann op <> ann rhs) lhs rhs
+    orf op lhs rhs = Term.or (ann lhs <> ann op <> ann rhs) lhs rhs
     and = andf <$> label "and" (reserved "&&")
-    andf op lhs rhs = Term.and (ann op <> ann rhs) lhs rhs
+    andf op lhs rhs = Term.and (ann lhs <> ann op <> ann rhs) lhs rhs
     infixApp = infixAppf <$> label "infixApp" (hashQualifiedInfixTerm <* optional semi)
-    infixAppf op lhs rhs = Term.apps op [(ann lhs, lhs), (ann rhs, rhs)]
+    infixAppf op lhs rhs = Term.apps' op [lhs, rhs]
 
 typedecl :: Var v => P v (L.Token v, Type v Ann)
 typedecl =

unison-cli/src/Unison/LSP/Queries.hs

@@ -0,0 +1,146 @@
+-- | Rewrites of some codebase queries, but which check the scratch file for info first.
+module Unison.LSP.Queries
+  ( refInTerm,
+    refInType,
+    findSmallestEnclosingNode,
+    findSmallestEnclosingType,
+    refInDecl,
+  )
+where
+
+import Control.Lens
+import qualified Unison.ABT as ABT
+import qualified Unison.ConstructorType as CT
+import qualified Unison.DataDeclaration as DD
+import Unison.LSP.Orphans ()
+import Unison.LabeledDependency
+import qualified Unison.LabeledDependency as LD
+import Unison.Lexer.Pos (Pos (..))
+import Unison.Parser.Ann (Ann)
+import qualified Unison.Parser.Ann as Ann
+import Unison.Prelude
+import Unison.Reference (TypeReference)
+import Unison.Symbol (Symbol)
+import Unison.Term (MatchCase (MatchCase), Term)
+import qualified Unison.Term as Term
+import Unison.Type (Type)
+import qualified Unison.Type as Type
+
+-- | Returns the reference a given term node refers to, if any.
+refInTerm :: (Term v a -> Maybe LabeledDependency)
+refInTerm term =
+  case ABT.out term of
+    ABT.Tm f -> case f of
+      Term.Int {} -> Nothing
+      Term.Nat {} -> Nothing
+      Term.Float {} -> Nothing
+      Term.Boolean {} -> Nothing
+      Term.Text {} -> Nothing
+      Term.Char {} -> Nothing
+      Term.Blank {} -> Nothing
+      Term.Ref ref -> Just (LD.TermReference ref)
+      Term.Constructor conRef -> Just (LD.ConReference conRef CT.Data)
+      Term.Request conRef -> Just (LD.ConReference conRef CT.Effect)
+      Term.Handle _a _b -> Nothing
+      Term.App _a _b -> Nothing
+      Term.Ann _a _typ -> Nothing
+      Term.List _xs -> Nothing
+      Term.If _cond _a _b -> Nothing
+      Term.And _l _r -> Nothing
+      Term.Or _l _r -> Nothing
+      Term.Lam _a -> Nothing
+      Term.LetRec _isTop _xs _y -> Nothing
+      Term.Let _isTop _a _b -> Nothing
+      Term.Match _a _cases -> Nothing
+      Term.TermLink ref -> Just (LD.TermReferent ref)
+      Term.TypeLink ref -> Just (LD.TypeReference ref)
+    ABT.Var _v -> Nothing
+    ABT.Cycle _r -> Nothing
+    ABT.Abs _v _r -> Nothing
+
+-- Returns the reference a given type node refers to, if any.
+refInType :: Type v a -> Maybe TypeReference
+refInType typ = case ABT.out typ of
+  ABT.Tm f -> case f of
+    Type.Ref ref -> Just ref
+    Type.Arrow _a _b -> Nothing
+    Type.Effect _a _b -> Nothing
+    Type.App _a _b -> Nothing
+    Type.Forall _r -> Nothing
+    Type.Ann _a _kind -> Nothing
+    Type.Effects _es -> Nothing
+    Type.IntroOuter _a -> Nothing
+  ABT.Var _v -> Nothing
+  ABT.Cycle _r -> Nothing
+  ABT.Abs _v _r -> Nothing
+
+-- | Find the the node in a term which contains the specified position, but none of its
+-- children contain that position.
+findSmallestEnclosingNode :: Pos -> Term Symbol Ann -> Maybe (Either (Term Symbol Ann) (Type Symbol Ann))
+findSmallestEnclosingNode pos term
+  | not (ABT.annotation term `Ann.contains` pos) = Nothing
+  | otherwise = (<|> Just (Left term)) $ do
+      case ABT.out term of
+        ABT.Tm f -> case f of
+          Term.Int {} -> Just (Left term)
+          Term.Nat {} -> Just (Left term)
+          Term.Float {} -> Just (Left term)
+          Term.Boolean {} -> Just (Left term)
+          Term.Text {} -> Just (Left term)
+          Term.Char {} -> Just (Left term)
+          Term.Blank {} -> Just (Left term)
+          Term.Ref {} -> Just (Left term)
+          Term.Constructor {} -> Just (Left term)
+          Term.Request {} -> Just (Left term)
+          Term.Handle a b -> findSmallestEnclosingNode pos a <|> findSmallestEnclosingNode pos b
+          Term.App a b -> findSmallestEnclosingNode pos a <|> findSmallestEnclosingNode pos b
+          Term.Ann a typ -> findSmallestEnclosingNode pos a <|> (Right <$> findSmallestEnclosingType pos typ)
+          Term.List xs -> altSum (findSmallestEnclosingNode pos <$> xs)
+          Term.If cond a b -> findSmallestEnclosingNode pos cond <|> findSmallestEnclosingNode pos a <|> findSmallestEnclosingNode pos b
+          Term.And l r -> findSmallestEnclosingNode pos l <|> findSmallestEnclosingNode pos r
+          Term.Or l r -> findSmallestEnclosingNode pos l <|> findSmallestEnclosingNode pos r
+          Term.Lam a -> findSmallestEnclosingNode pos a
+          Term.LetRec _isTop xs y -> altSum (findSmallestEnclosingNode pos <$> xs) <|> findSmallestEnclosingNode pos y
+          Term.Let _isTop a b -> findSmallestEnclosingNode pos a <|> findSmallestEnclosingNode pos b
+          Term.Match a cases ->
+            findSmallestEnclosingNode pos a
+              <|> altSum (cases <&> \(MatchCase _pat grd body) -> altSum (findSmallestEnclosingNode pos <$> grd) <|> findSmallestEnclosingNode pos body)
+          Term.TermLink {} -> Just (Left term)
+          Term.TypeLink {} -> Just (Left term)
+        ABT.Var _v -> Just (Left term)
+        ABT.Cycle r -> findSmallestEnclosingNode pos r
+        ABT.Abs _v r -> findSmallestEnclosingNode pos r
+
+-- | Find the the node in a type which contains the specified position, but none of its
+-- children contain that position.
+-- This is helpful for finding the specific type reference of a given argument within a type arrow
+-- that a position references.
+findSmallestEnclosingType :: Pos -> Type Symbol Ann -> Maybe (Type Symbol Ann)
+findSmallestEnclosingType pos typ
+  | not (ABT.annotation typ `Ann.contains` pos) = Nothing
+  | otherwise = (<|> Just typ) $ do
+      case ABT.out typ of
+        ABT.Tm f -> case f of
+          Type.Ref {} -> Just typ
+          Type.Arrow a b -> findSmallestEnclosingType pos a <|> findSmallestEnclosingType pos b
+          Type.Effect a b -> findSmallestEnclosingType pos a <|> findSmallestEnclosingType pos b
+          Type.App a b -> findSmallestEnclosingType pos a <|> findSmallestEnclosingType pos b
+          Type.Forall r -> findSmallestEnclosingType pos r
+          Type.Ann a _kind -> findSmallestEnclosingType pos a
+          Type.Effects es -> altSum (findSmallestEnclosingType pos <$> es)
+          Type.IntroOuter a -> findSmallestEnclosingType pos a
+        ABT.Var _v -> Just typ
+        ABT.Cycle r -> findSmallestEnclosingType pos r
+        ABT.Abs _v r -> findSmallestEnclosingType pos r
+
+-- | Returns the type reference the given position applies to within a Decl, if any.
+--
+-- I.e. if the cursor is over a type reference within a constructor signature or ability
+-- request signature, that type reference will be returned.
+refInDecl :: Pos -> DD.Decl Symbol Ann -> Maybe TypeReference
+refInDecl p (DD.asDataDecl -> dd) =
+  DD.constructors' dd
+    & altMap \(_conNameAnn, _v, typ) -> do
+      typeNode <- findSmallestEnclosingType p typ
+      ref <- refInType typeNode
+      pure ref

unison-cli/tests/Main.hs

@@ -7,13 +7,15 @@ import System.IO.CodePage (withCP65001)
 import qualified Unison.Test.ClearCache as ClearCache
 import qualified Unison.Test.Cli.Monad as Cli.Monad
 import qualified Unison.Test.GitSync as GitSync
+import qualified Unison.Test.LSP as LSP
 import qualified Unison.Test.UriParser as UriParser
 import qualified Unison.Test.VersionParser as VersionParser
 
 test :: Test ()
 test =
   tests
-    [ ClearCache.test,
+    [ LSP.test,
+      ClearCache.test,
       Cli.Monad.test,
       GitSync.test,
       UriParser.test,

unison-cli/tests/Unison/Test/LSP.hs

@@ -0,0 +1,106 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE QuasiQuotes #-}
+
+module Unison.Test.LSP (test) where
+
+import qualified Crypto.Random as Random
+import Data.Bifunctor (bimap)
+import Data.List.Extra (firstJust)
+import Data.String.Here.Uninterpolated (here)
+import Data.Text
+import qualified Data.Text as Text
+import EasyTest
+import qualified System.IO.Temp as Temp
+import qualified Unison.ABT as ABT
+import qualified Unison.Cli.TypeCheck as Typecheck
+import Unison.Codebase (Codebase)
+import qualified Unison.Codebase.Init as Codebase.Init
+import qualified Unison.Codebase.SqliteCodebase as SC
+import qualified Unison.LSP.Queries as LSPQ
+import qualified Unison.Lexer.Pos as Lexer
+import Unison.Parser.Ann (Ann (..))
+import Unison.Prelude
+import qualified Unison.Result as Result
+import Unison.Symbol (Symbol)
+import qualified Unison.Syntax.Lexer as L
+import qualified Unison.Syntax.Parser as Parser
+import Unison.Term (Term)
+import qualified Unison.Term as Term
+import Unison.Type (Type)
+import qualified Unison.Type as Type
+import qualified Unison.UnisonFile as UF
+
+test :: Test ()
+test =
+  scope "annotations" . tests . fmap makeNodeSelectionTest $
+    [ ( "Binary Op lhs",
+        [here|term = tr|ue && false|],
+        True,
+        Left (Term.Boolean True)
+      ),
+      ( "Binary Op rhs",
+        [here|term = true && fa|lse|],
+        True,
+        Left (Term.Boolean False)
+      ),
+      ( "Custom Op lhs",
+        [here|
+a &&& b = a && b
+term = tr|ue &&& false
+|],
+        True,
+        Left (Term.Boolean True)
+      )
+    ]
+
+-- | Test helper which lets you specify a cursor position inline with source text as a '|'.
+extractCursor :: Text -> Test (Lexer.Pos, Text)
+extractCursor txt =
+  case Text.splitOn "|" txt of
+    [before, after] ->
+      let col = Text.length $ Text.takeWhileEnd (/= '\n') before
+          line = Prelude.length $ Text.lines before
+       in pure $ (Lexer.Pos line col, before <> after)
+    _ -> crash "expected exactly one cursor"
+
+makeNodeSelectionTest :: (String, Text, Bool, Either ((Term.F Symbol Ann Ann (Term Symbol Ann))) (Type.F (Type Symbol Ann))) -> Test ()
+makeNodeSelectionTest (name, testSrc, testTypechecked, expected) = scope name $ do
+  (pos, src) <- extractCursor testSrc
+  (mayParsedFile, mayTypecheckedFile) <- withTestCodebase \codebase -> do
+    let generateUniqueName = Parser.uniqueBase32Namegen <$> Random.getSystemDRG
+    let ambientAbilities = []
+    let parseNames = mempty
+    let lexedSource = (src, L.lexer name (Text.unpack src))
+    r <- Typecheck.typecheckHelper codebase generateUniqueName ambientAbilities parseNames (Text.pack name) lexedSource
+    let Result.Result _notes mayResult = r
+    let (parsedFile, typecheckedFile) = case mayResult of
+          Nothing -> (Nothing, Nothing)
+          Just (Left uf) -> (Just uf, Nothing)
+          Just (Right tf) -> (Just $ UF.discardTypes tf, Just tf)
+    pure (parsedFile, typecheckedFile)
+  scope "parsed file" $ do
+    pf <- maybe (crash "Failed to parse") pure mayParsedFile
+    let pfResult =
+          UF.terms pf
+            & firstJust \(_v, trm) ->
+              LSPQ.findSmallestEnclosingNode pos trm
+    expectEqual (Just $ bimap ABT.Tm ABT.Tm expected) (bimap ABT.out ABT.out <$> pfResult)
+
+  when testTypechecked $
+    scope "typechecked file" $ do
+      tf <- maybe (crash "Failed to typecheck") pure mayTypecheckedFile
+      let tfResult =
+            UF.hashTermsId tf
+              & toList
+              & firstJust \(_refId, _wk, trm, _typ) ->
+                LSPQ.findSmallestEnclosingNode pos trm
+      expectEqual (Just $ bimap ABT.Tm ABT.Tm expected) (bimap ABT.out ABT.out <$> tfResult)
+
+withTestCodebase ::
+  (Codebase IO Symbol Ann -> IO r) -> Test r
+withTestCodebase action = do
+  r <- io do
+    tmp <- Temp.getCanonicalTemporaryDirectory
+    tmpDir <- Temp.createTempDirectory tmp "lsp-test"
+    Codebase.Init.withCreatedCodebase SC.init "lsp-test" tmpDir SC.DontLock action
+  either (crash . show) pure r

unison-cli/unison-cli.cabal

@@ -83,6 +83,7 @@ library
       Unison.LSP.Hover
       Unison.LSP.NotificationHandlers
       Unison.LSP.Orphans
+      Unison.LSP.Queries
       Unison.LSP.Types
       Unison.LSP.UCMWorker
       Unison.LSP.VFS
@@ -580,6 +581,7 @@ test-suite cli-tests
       Unison.Test.ClearCache
       Unison.Test.Cli.Monad
       Unison.Test.GitSync
+      Unison.Test.LSP
       Unison.Test.Ucm
       Unison.Test.UriParser
       Unison.Test.VersionParser

unison-syntax/src/Unison/Parser/Ann.hs

@@ -20,9 +20,28 @@ instance Monoid Ann where
   mempty = External
 
 instance Semigroup Ann where
-  Ann s1 _ <> Ann _ e2 = Ann s1 e2
+  Ann s1 e1 <> Ann s2 e2 = Ann (min s1 s2) (max e1 e2)
   -- If we have a concrete location from a file, use it
   External <> a = a
   a <> External = a
   Intrinsic <> a = a
   a <> Intrinsic = a
+
+-- | Checks whether an annotation contains a given position
+-- i.e. pos ∈ [start, end)
+--
+-- >>> Intrinsic `contains` L.Pos 1 1
+-- False
+--
+-- >>> External `contains` L.Pos 1 1
+-- False
+--
+-- >>> Ann (L.Pos 0 0) (L.Pos 0 10) `contains` L.Pos 0 5
+-- True
+--
+-- >>> Ann (L.Pos 0 0) (L.Pos 0 10) `contains` L.Pos 0 10
+-- False
+contains :: Ann -> L.Pos -> Bool
+contains Intrinsic _ = False
+contains External _ = False
+contains (Ann start end) p = start <= p && p < end