Fix letrec behaviour

tools/lsp/squirrel/src/AST/Parser.hs

@@ -11,6 +11,8 @@ import Data.Sum
 import AST.Types
 
 import Parser
+import Range
+import Product
 import Tree hiding (skip)
 
 -- import Debug.Trace
@@ -29,13 +31,31 @@ ranged p = do
 -- | The entrypoint.
 contract :: Parser (Pascal ASTInfo)
 contract =
-  ranged do
-    pure Contract
+    pure contract'
+      <*> getInfo
       <*> subtree "contract" do
             many do
               inside "declaration:" do
                 declaration
 
+  where
+    contract'
+      :: ASTInfo
+      -> [Pascal ASTInfo]
+      -> Pascal ASTInfo
+    contract' r = foldr (contract'' $ getElem r) (mk r ContractEnd)
+
+    contract''
+      :: Range
+      -> Pascal ASTInfo
+      -> Pascal ASTInfo
+      -> Pascal ASTInfo
+    contract'' r x xs = mk (Cons r' rest) $ ContractCons x xs
+      where
+        r' = Range start end f
+        Range _ end f = r
+        Cons (Range start _ _) rest = infoOf x
+
 name :: Parser (Pascal ASTInfo)
 name = ranged do pure Name <*> token "Name"
 
@@ -745,10 +765,26 @@ letExpr = do
               declaration <|> statement
       <*> inside "body"expr
   where
+    let'
+      :: ASTInfo
+      -> (Maybe [Pascal ASTInfo])
+      -> Pascal ASTInfo
+      -> Pascal ASTInfo
     let' r decls body = case decls of
-      Just them -> mk r $ Let them body
+      Just them -> foldr (let'' $ getElem r) body them
       Nothing   -> body
 
+    let''
+      :: Range
+      -> Pascal ASTInfo
+      -> Pascal ASTInfo
+      -> Pascal ASTInfo
+    let'' r decl b = mk (Cons r' rest) $ Let decl b
+      where
+        r' = Range start end f
+        Range _ end f = r
+        Cons (Range start _ _) rest = infoOf b
+
 statement :: Parser (Pascal ASTInfo)
 statement = ranged do pure Action <*> expr
 

tools/lsp/squirrel/src/AST/Types.hs

@@ -24,7 +24,8 @@ type Pascal = Tree
   ]
 
 data Contract it
-  = Contract      [it] -- ^ Declaration
+  = ContractEnd
+  | ContractCons it it -- ^ Declaration
   deriving (Show) via PP (Contract it)
   deriving stock (Functor, Foldable, Traversable)
 
@@ -78,7 +79,7 @@ data TField it
 
 -- | TODO: break onto smaller types? Literals -> Constant; mapOps; mmove Annots to Decls.
 data Expr it
-  = Let       [it] it -- [Declaration] (Expr)
+  = Let       it it -- Declaration (Expr)
   | Apply     it [it] -- (Expr) [Expr]
   | Constant  it -- (Constant)
   | Ident     it -- (QualifiedName)
@@ -186,8 +187,8 @@ newtype FieldName it = FieldName Text
 
 instance Pretty1 Contract where
   pp1 = \case
-    Contract decls ->
-      sparseBlock decls
+    ContractEnd -> "(* end *)"
+    ContractCons x xs -> x $$ " " $$ xs
 
 instance Pretty1 Declaration where
   pp1 = \case
@@ -243,7 +244,7 @@ instance Pretty1 Variant where
 
 instance Pretty1 Expr where
   pp1 = \case
-    Let       decls body -> "block {" `indent` sparseBlock decls `above` "}" <+> "with" `indent` body
+    Let       decl body -> "block {" `indent` decl `above` "}" <+> "with" `indent` body
     Apply     f xs       -> f <+> tuple xs
     Constant  constant   -> constant
     Ident     qname      -> qname