From 6bc0965232a3bac127fc736dfa2cc4ffebcbbec9 Mon Sep 17 00:00:00 2001
From: Jinwoo Lee <jinwoo@google.com>
Date: Thu, 28 May 2020 10:45:45 -0700
Subject: [PATCH] Support arithmetic sequences.

`from`      : [a .. ]
`fromThen`  : [a, b .. ]
`fromTo`    : [a .. b]
`fromThenTo`: [a, b .. c]

These are needed before supporting list comprehensions (#64).
---
 .gitignore                |  1 +
 src/GHC/SourceGen/Expr.hs | 41 +++++++++++++++++++++++++++++++++++++++
 tests/pprint_test.hs      |  6 ++++++
 3 files changed, 48 insertions(+)

diff --git a/.gitignore b/.gitignore
index e9893a3..a86bea9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,3 +2,4 @@
 dist/
 # Generated automatically by hpack:
 *.cabal
+*.yaml.lock
diff --git a/src/GHC/SourceGen/Expr.hs b/src/GHC/SourceGen/Expr.hs
index 73a5585..a5e7dd9 100644
--- a/src/GHC/SourceGen/Expr.hs
+++ b/src/GHC/SourceGen/Expr.hs
@@ -21,6 +21,10 @@ module GHC.SourceGen.Expr
     , tyApp
     , recordConE
     , recordUpd
+    , from
+    , fromThen
+    , fromTo
+    , fromThenTo
     ) where
 
 import GHC.Hs.Expr
@@ -187,3 +191,40 @@ recordUpd e fs =
             }
     withPlaceHolder4 = withPlaceHolder . withPlaceHolder . withPlaceHolder
                             . withPlaceHolder
+
+-- | An arithmetic sequence expression with a start value.
+--
+-- > [a ..]
+-- > =====
+-- > from (var "a")
+from :: HsExpr' -> HsExpr'
+from from' =
+    noExt ArithSeq Nothing $ From (builtLoc from')
+
+-- | An arithmetic sequence expression with a start and a step values.
+--
+-- > [a, b ..]
+-- > =====
+-- > fromThen (var "a") (var "b")
+fromThen :: HsExpr' -> HsExpr' -> HsExpr'
+fromThen from' then' =
+    noExt ArithSeq Nothing $ FromThen (builtLoc from') (builtLoc then')
+
+-- | An arithmetic sequence expression with a start and an end values.
+--
+-- > [a .. b]
+-- > =====
+-- > fromTo (var "a") (var "b")
+fromTo :: HsExpr' -> HsExpr' -> HsExpr'
+fromTo from' to =
+    noExt ArithSeq Nothing $ FromTo (builtLoc from') (builtLoc to)
+
+-- | An arithmetic sequence expression with a start, a step, and an end values.
+--
+-- > [a, b .. c]
+-- > =====
+-- > fromThenTo (var "a") (var "b") (var "c")
+fromThenTo :: HsExpr' -> HsExpr' -> HsExpr' -> HsExpr'
+fromThenTo from' then' to =
+    noExt ArithSeq Nothing $
+        FromThenTo (builtLoc from') (builtLoc then') (builtLoc to)
diff --git a/tests/pprint_test.hs b/tests/pprint_test.hs
index 00c91c8..4fc6770 100644
--- a/tests/pprint_test.hs
+++ b/tests/pprint_test.hs
@@ -221,6 +221,12 @@ exprsTest dflags = testGroup "Expr"
         -- TODO: add more tests.
         [ "do (let x = 1 in x)" :~ do' [stmt $ let' [valBind "x" (int 1)] (var "x")]
         ]
+    , test "arithSeq"
+        [ "[a .. ]" :~ from (var "a")
+        , "[a, b .. ]" :~ fromThen (var "a") (var "b")
+        , "[a .. b]" :~ fromTo (var "a") (var "b")
+        , "[a, b .. c]" :~ fromThenTo (var "a") (var "b") (var "c")
+        ]
     ]
   where
     test = testExprs dflags