From 72df0644326582ff828dc7be558386cbc24a461a Mon Sep 17 00:00:00 2001
From: Gabriel Gonzalez <Gabriel439@gmail.com>
Date: Tue, 1 Jun 2021 17:47:39 -0700
Subject: [PATCH] Sort `let` bindings with imports

Fixes https://github.com/dhall-lang/dhall-haskell/issues/2204

This sorts imports within a sequence of `let` bindings as high up as they
will go, without reordering `let` bindings of the same name.
---
 dhall/src/Dhall/Lint.hs                  | 63 ++++++++++++++++++++++--
 dhall/tests/lint/success/sortLetsA.dhall | 11 +++++
 dhall/tests/lint/success/sortLetsB.dhall | 11 +++++
 3 files changed, 80 insertions(+), 5 deletions(-)
 create mode 100644 dhall/tests/lint/success/sortLetsA.dhall
 create mode 100644 dhall/tests/lint/success/sortLetsB.dhall

diff --git a/dhall/src/Dhall/Lint.hs b/dhall/src/Dhall/Lint.hs
index 4c75c7ac0..e7bc4b082 100644
--- a/dhall/src/Dhall/Lint.hs
+++ b/dhall/src/Dhall/Lint.hs
@@ -19,6 +19,7 @@ module Dhall.Lint
     ) where
 
 import Control.Applicative ((<|>))
+import Data.List.NonEmpty (NonEmpty(..))
 
 import Dhall.Syntax
     ( Binding (..)
@@ -37,9 +38,10 @@ import Dhall.Syntax
 
 import qualified Data.Foldable      as Foldable
 import qualified Data.List.NonEmpty as NonEmpty
+import qualified Data.Map           as Map
 import qualified Data.Text          as Text
 import qualified Dhall.Core         as Core
-import qualified Dhall.Map          as Map
+import qualified Dhall.Map
 import qualified Dhall.Optics
 import qualified Lens.Family
 
@@ -52,7 +54,7 @@ import qualified Lens.Family
     * consolidates nested @let@ bindings to use a multiple-@let@ binding with 'removeLetInLet'
     * fixes paths of the form @.\/..\/foo@ to @..\/foo@
 -}
-lint :: Expr s Import -> Expr s Import
+lint :: Eq s => Expr s Import -> Expr s Import
 lint =  Dhall.Optics.rewriteOf subExpressions rewrite
   where
     rewrite e =
@@ -61,6 +63,7 @@ lint =  Dhall.Optics.rewriteOf subExpressions rewrite
         <|> fixParentPath            e
         <|> removeLetInLet           e
         <|> addPreludeExtensions     e
+        <|> sortImports              e
 
 -- | Remove unused `Let` bindings.
 removeUnusedBindings :: Eq a => Expr s a -> Maybe (Expr s a)
@@ -176,7 +179,7 @@ useToMap
             (Core.shallowDenote -> App
                 (Core.shallowDenote -> List)
                 (Core.shallowDenote -> Record
-                    (Map.sort ->
+                    (Dhall.Map.sort ->
                         [ ("mapKey", Core.shallowDenote . Core.recordFieldValue -> Text)
                         , ("mapValue", _)
                         ]
@@ -192,14 +195,14 @@ useToMap (ListLit _ keyValues)
     , Just keyValues' <- traverse convert keyValues =
         Just
             (ToMap
-                (RecordLit (Map.fromList (Foldable.toList keyValues')))
+                (RecordLit (Dhall.Map.fromList (Foldable.toList keyValues')))
                 Nothing
             )
   where
     convert keyValue =
         case Core.shallowDenote keyValue of
             RecordLit
-                (Map.sort ->
+                (Dhall.Map.sort ->
                     [ ("mapKey"  , Core.shallowDenote . Core.recordFieldValue -> TextLit (Chunks [] key))
                     , ("mapValue", value)
                     ]
@@ -209,3 +212,53 @@ useToMap (ListLit _ keyValues)
                 Nothing
 useToMap _ =
     Nothing
+
+-- | This sorts `let` bindings to move imports to the front if doing so does not
+-- change the behavior of the code.
+sortImports :: Eq s => Expr s Import -> Maybe (Expr s Import)
+sortImports oldExpression@(Let binding0 oldBody0)
+    | oldExpression == newExpression = Nothing
+    | otherwise                      = Just newExpression
+  where
+    toBool (Embed _ ) = False
+    toBool (Note _ e) = toBool e
+    toBool  _         = True
+
+    process (seen, index) Binding{..} oldBody function = (pair, pairs, newBody)
+      where
+        order =
+            if b then index else Map.findWithDefault (0 :: Int) variable seen
+
+        b = toBool value
+
+        pair = (order, function)
+
+        ~(pairs, newBody) =
+            label (Map.insert variable order seen, index + 1) oldBody
+
+    label state (Let binding oldBody) = (pair : pairs, newBody)
+      where
+        function = Let binding
+
+        ~(pair, pairs, newBody) = process state binding oldBody function
+
+    label state (Note src (Let binding oldBody)) = (pair : pairs, newBody)
+      where
+        function e = Note src (Let binding e)
+
+        ~(pair, pairs, newBody) = process state binding oldBody function
+
+    label _ body =
+        ([], body)
+
+    ~(pairs0, newBody0) = (pair :| pairs, newBody)
+      where
+        function = Let binding0
+
+        ~(pair, pairs, newBody) =
+            process (Map.empty, 1) binding0 oldBody0 function
+
+    sortedFunctions = fmap snd (NonEmpty.sortWith fst pairs0)
+
+    newExpression = foldr id newBody0 sortedFunctions
+sortImports _ = Nothing
diff --git a/dhall/tests/lint/success/sortLetsA.dhall b/dhall/tests/lint/success/sortLetsA.dhall
new file mode 100644
index 000000000..5abd879ae
--- /dev/null
+++ b/dhall/tests/lint/success/sortLetsA.dhall
@@ -0,0 +1,11 @@
+let a = 1
+
+let b = 2
+
+let a = https://example.com
+
+let c = https://example.com
+
+let d = c
+
+in  [ a@1, a, b, c, d ]
diff --git a/dhall/tests/lint/success/sortLetsB.dhall b/dhall/tests/lint/success/sortLetsB.dhall
new file mode 100644
index 000000000..9e7f11bbf
--- /dev/null
+++ b/dhall/tests/lint/success/sortLetsB.dhall
@@ -0,0 +1,11 @@
+let c = https://example.com/
+
+let a = 1
+
+let a = https://example.com/
+
+let b = 2
+
+let d = c
+
+in  [ a@1, a, b, c, d ]