Add required properties and refinement types to traits

Author: rockofox

app/Main.hs

@@ -7,7 +7,7 @@ import AST qualified
 import BytecodeCompiler (CompilerError (..), compileFail, renderCompilerErrors)
 import BytecodeCompiler qualified
 import Control.Exception
-import Control.Monad (filterM)
+import Control.Monad (filterM, when)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Control.Monad.Identity (Identity (..))
 import Control.Monad.State (StateT (runStateT), evalStateT)

lib/AST.hs

@@ -57,7 +57,7 @@ data Expr
     | Cast {castExpr :: Expr, castType :: Expr, castPos :: Position}
     | TypeLit {typeLitType :: Type, typeLitPos :: Position}
     | Flexible {flexibleExpr :: Expr, flexiblePos :: Position}
-    | Trait {name :: String, methods :: [Expr], generics :: [GenericExpr], traitPos :: Position}
+    | Trait {name :: String, methods :: [Expr], generics :: [GenericExpr], requiredProperties :: [(String, Type)], refinement :: Maybe Expr, refinementSrc :: String, traitPos :: Position}
     | Impl {trait :: String, traitTypeArgs :: [Type], for :: Type, methods :: [Expr], implPos :: Position}
     | StrictEval {strictEvalExpr :: Expr, strictEvalPos :: Position}
     | External {externalName :: String, externalArgs :: [Expr], externalPos :: Position}
@@ -119,7 +119,7 @@ children (Lambda _ a _) = [a]
 children (Cast a b _) = [a, b]
 children (TypeLit _ _) = []
 children (Flexible a _) = [a]
-children (Trait _ a _ _) = a
+children (Trait _ a _ _ r _ _) = a ++ maybeToList r
 children (Impl _ _ _ a _) = a
 children (FuncDec{}) = []
 children (StrictEval a _) = [a]

lib/BytecodeCompiler.hs

@@ -770,7 +770,7 @@ parseFreeUnsafe t = case parseProgram t initCompilerFlags{needsMain = False} of
     replacePositionWithAnyPosition (Cast a b _) = Cast (replacePositionWithAnyPosition a) (replacePositionWithAnyPosition b) anyPosition
     replacePositionWithAnyPosition (TypeLit type' _) = TypeLit type' anyPosition
     replacePositionWithAnyPosition (Flexible a _) = Flexible (replacePositionWithAnyPosition a) anyPosition
-    replacePositionWithAnyPosition (Trait name methods g _) = Trait name (map replacePositionWithAnyPosition methods) g anyPosition
+    replacePositionWithAnyPosition (Trait name methods g reqProps ref refSrc _) = Trait name (map replacePositionWithAnyPosition methods) g reqProps (fmap replacePositionWithAnyPosition ref) refSrc anyPosition
     replacePositionWithAnyPosition (Impl traitName traitTypeArgs f m _) = Impl traitName traitTypeArgs f (map replacePositionWithAnyPosition m) anyPosition
     replacePositionWithAnyPosition (StrictEval a _) = StrictEval (replacePositionWithAnyPosition a) anyPosition
     replacePositionWithAnyPosition (External n a _) = External n (map replacePositionWithAnyPosition a) anyPosition
@@ -867,18 +867,57 @@ trait = do
     keyword "trait"
     name <- identifier <?> "trait name"
     generics <- fromMaybe [] <$> optional generic <?> "trait generics"
-    methods <-
-        ( do
+    hasEquals <- lookAhead (optional (try (symbol "=")))
+    (requiredProperties, refinementSrc, refinement, methods) <- case hasEquals of
+        Just _ -> do
             symbol "="
-            keyword "do"
-            newline'
-            fds <- funcDec `sepEndBy` newline' <?> "trait methods"
-            keyword "end"
-            return fds
-        )
-            <|> return []
+            reqProps <- optional (try (parens (traitField `sepBy` symbol ",")) <?> "required properties")
+            refSrc <- lookAhead $ optional $ do
+                keyword "satisfies"
+                parens (many (noneOf [')'])) <?> "refinement source"
+            ref <- optional $ do
+                keyword "satisfies"
+                parens expr <?> "refinement"
+            methods' <-
+                ( do
+                    keyword "do"
+                    newline'
+                    fds <- funcDec `sepEndBy` newline' <?> "trait methods"
+                    keyword "end"
+                    return fds
+                )
+                    <|> return []
+            return (reqProps, refSrc, ref, methods')
+        Nothing -> do
+            reqProps <- optional (try (parens (traitField `sepBy` symbol ",")) <?> "required properties")
+            refSrc <- lookAhead $ optional $ do
+                keyword "satisfies"
+                parens (many (noneOf [')'])) <?> "refinement source"
+            ref <- optional $ do
+                keyword "satisfies"
+                parens expr <?> "refinement"
+            hasEqualsAfter <- lookAhead (optional (try (symbol "=")))
+            methods' <- case hasEqualsAfter of
+                Just _ -> do
+                    symbol "="
+                    ( do
+                            keyword "do"
+                            newline'
+                            fds <- funcDec `sepEndBy` newline' <?> "trait methods"
+                            keyword "end"
+                            return fds
+                        )
+                        <|> return []
+                Nothing -> return []
+            return (reqProps, refSrc, ref, methods')
     end <- getOffset
-    return $ Trait name methods generics (Position (start, end))
+    return $ Trait name methods generics (fromMaybe [] requiredProperties) refinement (fromMaybe "" refinementSrc) (Position (start, end))
+  where
+    traitField = do
+        fieldName <- identifier <?> "field name"
+        symbol ":"
+        fieldType <- validType <?> "field type"
+        return (fieldName, fieldType)
 
 impl :: Parser Expr
 impl = do

lib/Parser.hs

@@ -1,3 +1,9 @@
 import qualified Module2
 
 let x = Module2.
+
+
+
+
+
+

lsp/tests/data/completion_imports.in

@@ -1,3 +1,9 @@
 import Module2
 
 let x = 
+
+
+
+
+
+

lsp/tests/data/completion_imports_unqualified.in

@@ -1,3 +1,9 @@
 import NonExistentModule
 
 let x = 42
+
+
+
+
+
+

lsp/tests/data/diagnostics_missing_module.in

@@ -1,3 +1,9 @@
 import qualified Module2
 
 let x = Module2.add
+
+
+
+
+
+

lsp/tests/data/hover_qualified.in

@@ -5,3 +5,9 @@ let add (a: Int b: Int) : Int = a + b
 let multiply (a: Int b: Int) : Int = a * b
 
 let greet (name: String) : String = "Hello, " ++ name ++ "!"
+
+
+
+
+
+