From 9f29dab195165f540be04cda814f6ecf01430a19 Mon Sep 17 00:00:00 2001
From: Kirill Andreev <hindmost.one@gmail.com>
Date: Tue, 28 Jul 2020 20:51:32 +0400
Subject: [PATCH] Fix letrec partially

---
 src/test/contracts/closure-3.ligo   |  4 ++--
 tools/lsp/squirrel/src/AST/Scope.hs | 36 +++++++++++++++++++++++++----
 tools/lsp/squirrel/src/AST/Types.hs |  2 +-
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/src/test/contracts/closure-3.ligo b/src/test/contracts/closure-3.ligo
index ae4e68956..b860c8593 100644
--- a/src/test/contracts/closure-3.ligo
+++ b/src/test/contracts/closure-3.ligo
@@ -5,6 +5,6 @@
 function foobar (const i : int) : int is
   block {
     const j : int = 3;
-    const k : int = 4;
-    function add (const l : int) : int is i+j+k+l
+    const k : int = j;
+    const l : int = l;
   } with add (42)
diff --git a/tools/lsp/squirrel/src/AST/Scope.hs b/tools/lsp/squirrel/src/AST/Scope.hs
index 67ef19116..11b7b6418 100644
--- a/tools/lsp/squirrel/src/AST/Scope.hs
+++ b/tools/lsp/squirrel/src/AST/Scope.hs
@@ -15,6 +15,7 @@ module AST.Scope
 
 import           Control.Arrow (first, second)
 import           Control.Monad.State
+import           Control.Monad.Identity
 
 import qualified Data.List   as List
 import           Data.Map            (Map)
@@ -91,8 +92,9 @@ ofCategory Type     ScopedDecl { _sdType = Just (Right Star) } = True
 ofCategory _        _                                          = False
 
 instance Modifies (Product '[[ScopedDecl], Maybe Category, [Text], Range, a]) where
-  ascribe (ds :> _) d =
-    color 2 (fsep (map (pp . _sdName) ds))
+  ascribe (ds :> _ :> _ :> r :> _) d =
+    color 3 (fsep (map (pp . _sdName) ds))
+    $$ pp r
     $$ d
 
 addLocalScopes
@@ -102,8 +104,34 @@ addLocalScopes
 addLocalScopes tree =
     fmap (\xs -> fullEnvAt envWithREfs (getRange xs) :> xs) tree1
   where
-    tree1       = addNameCategories tree
-    envWithREfs = getEnvTree tree
+    tree0       = runIdentity $ unLetRec tree
+    tree1       = addNameCategories tree0
+    envWithREfs = getEnvTree tree0
+
+unLetRec
+  :: forall xs m
+  .  ( Monad m
+     , Contains Range xs
+     , Eq (Product xs)
+     )
+  => LIGO (Product xs)
+  -> m (LIGO (Product xs))
+unLetRec = descent leaveBe
+  [ Descent
+    [ \case
+       (r, Let (layer -> Just (Seq xs)) b) -> return $ convert (getElem r) b xs
+       _                                   -> return Nothing
+    ]
+  ]
+  where
+    convert :: Range -> LIGO (Product xs) -> [LIGO (Product xs)] -> Maybe (Product xs, Expr (LIGO (Product xs)))
+    convert r b = match @Expr . linearize r b
+
+    linearize :: Range -> LIGO (Product xs) -> [LIGO (Product xs)] -> LIGO (Product xs)
+    linearize r b [x]      = make (modElem @Range (delta r) $ extract x, Let x b)
+    linearize r b (x : xs) = make (modElem @Range (delta r) $ extract x, Let x (linearize r b xs))
+
+    delta (Range _ f _) (Range s _ t) = Range s f t
 
 addNameCategories
   :: (Contains Range xs, Eq (Product xs))
diff --git a/tools/lsp/squirrel/src/AST/Types.hs b/tools/lsp/squirrel/src/AST/Types.hs
index 8d05be0e9..c2ce71a6f 100644
--- a/tools/lsp/squirrel/src/AST/Types.hs
+++ b/tools/lsp/squirrel/src/AST/Types.hs
@@ -288,7 +288,7 @@ instance Pretty1 Variant where
 
 instance Pretty1 Expr where
   pp1 = \case
-    Let       decl body  -> decl `above` "with" `indent` body
+    Let       decl body  -> "let" <+> decl `above` body
     Apply     f xs       -> f <+> xs
     Constant  constant   -> constant
     Ident     qname      -> qname