Make function chains pretty (#1479)

- closes #1473

scripts/compare-format.sh

@@ -0,0 +1,80 @@
+#!/bin/bash
+
+MIN=10
+MAX=200
+STEP=5
+
+function help() {
+    echo "$0 [--min $MIN --max $MAX --step $STEP] EXAMPLE_FILE.sw"
+    echo
+    echo "This script helps to compare the format layout for a range"
+    echo "of output widths. Afterwards it prints those that differ"
+    echo "as markdown with triple backticks."
+}
+
+# Simple argument parsing from https://stackoverflow.com/a/14203146
+POSITIONAL_ARGS=()
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --min)
+      MIN="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    --max)
+      MAX="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -s|--step)
+      STEP="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -h|--help)
+      help
+      exit 0
+      ;;
+    -*)
+      echo "Unknown option $1"
+      help
+      exit 1
+      ;;
+    *)
+      POSITIONAL_ARGS+=("$1") # save positional arg
+      shift # past argument
+      ;;
+  esac
+done
+
+# Build first, otherwise the ouput would go to temporary files
+cabal build -O0
+
+function compare_format() {
+    echo "# $1"
+    if ! test -f "$1"; then
+      echo "Could not find file '$1'!"
+      return
+    fi
+    # save each version in a temporary file
+    t=$(mktemp -d)
+
+    for i in $(seq "$MIN" "$STEP" "$MAX"); do
+        echo -en "${i}\r"  # show progress
+        cabal run swarm -O0 -- format --width "$i" "$1" > "$t/$i.sw";
+    done
+    echo -en "   \r"
+
+    for i in $(seq "$MAX" "-$STEP" "$MIN"); do
+    if ! cmp -s "$t/$i.sw" "$t"/$((i - STEP)).sw; then
+        echo "$i";
+        echo '```';
+        cat "$t/$i.sw";
+        echo '```';
+    fi;
+    done
+}
+
+for file in "${POSITIONAL_ARGS[@]}"; do
+  compare_format "$file";
+done

src/Swarm/Language/Pretty.hs

@@ -30,7 +30,7 @@ import Swarm.Language.Parse (getLocRange)
 import Swarm.Language.Syntax
 import Swarm.Language.Typecheck
 import Swarm.Language.Types
-import Swarm.Util (showEnum, showLowT)
+import Swarm.Util (showEnum, showLowT, unsnocNE)
 import Witch
 
 ------------------------------------------------------------
@@ -126,6 +126,19 @@ data Wildcard = Wildcard
 instance PrettyPrec Wildcard where
   prettyPrec _ _ = "_"
 
+-- | Split a function type chain, so that we can pretty print
+--   the type parameters aligned on each line when they don't fit.
+class UnchainableFun t where
+  unchainFun :: t -> [t]
+
+instance UnchainableFun Type where
+  unchainFun (a :->: ty) = a : unchainFun ty
+  unchainFun ty = [ty]
+
+instance UnchainableFun (UTerm TypeF ty) where
+  unchainFun (UTerm (TyFunF ty1 ty2)) = ty1 : unchainFun ty2
+  unchainFun ty = [ty]
+
 instance (PrettyPrec (t (Fix t))) => PrettyPrec (Fix t) where
   prettyPrec p = prettyPrec p . unFix
 
@@ -133,7 +146,7 @@ instance (PrettyPrec (t (UTerm t v)), PrettyPrec v) => PrettyPrec (UTerm t v) wh
   prettyPrec p (UTerm t) = prettyPrec p t
   prettyPrec p (UVar v) = prettyPrec p v
 
-instance (PrettyPrec t) => PrettyPrec (TypeF t) where
+instance ((UnchainableFun t), (PrettyPrec t)) => PrettyPrec (TypeF t) where
   prettyPrec _ (TyBaseF b) = ppr b
   prettyPrec _ (TyVarF v) = pretty v
   prettyPrec p (TySumF ty1 ty2) =
@@ -145,8 +158,12 @@ instance (PrettyPrec t) => PrettyPrec (TypeF t) where
   prettyPrec p (TyCmdF ty) = pparens (p > 9) $ "cmd" <+> prettyPrec 10 ty
   prettyPrec _ (TyDelayF ty) = braces $ ppr ty
   prettyPrec p (TyFunF ty1 ty2) =
-    pparens (p > 0) $
-      prettyPrec 1 ty1 <+> "->" <+> prettyPrec 0 ty2
+    let (iniF, lastF) = unsnocNE $ ty1 NE.:| unchainFun ty2
+        funs = (prettyPrec 1 <$> iniF) <> [ppr lastF]
+        inLine l r = l <+> "->" <+> r
+        multiLine l r = l <+> "->" <> hardline <> r
+     in pparens (p > 0) . align $
+          flatAlt (concatWith multiLine funs) (concatWith inLine funs)
   prettyPrec _ (TyRcdF m) = brackets $ hsep (punctuate "," (map prettyBinding (M.assocs m)))
 
 instance PrettyPrec Polytype where
@@ -219,27 +236,15 @@ instance PrettyPrec Term where
             _ -> prettyPrecApp p t1 t2
     _ -> prettyPrecApp p t1 t2
   prettyPrec _ (TLet _ x mty t1 t2) =
-    group . vsep $
-      [ hsep $
-          ["let", pretty x]
-            ++ maybe [] (\ty -> [":", ppr ty]) mty
-            ++ ["=", ppr t1, "in"]
+    sep
+      [ prettyDefinition "let" x mty t1 <+> "in"
       , ppr t2
       ]
   prettyPrec _ (TDef _ x mty t1) =
-    let (t1rest, t1lams) = unchainLambdas t1
-     in group . vsep $
-          [ nest 2 $
-              vsep
-                [ hsep $
-                    ["def", pretty x]
-                      ++ maybe [] (\ty -> [":", ppr ty]) mty
-                      ++ ["="]
-                      ++ map prettyLambda t1lams
-                , ppr t1rest
-                ]
-          , "end"
-          ]
+    sep
+      [ prettyDefinition "def" x mty t1
+      , "end"
+      ]
   prettyPrec p (TBind Nothing t1 t2) =
     pparens (p > 0) $
       prettyPrec 1 t1 <> ";" <> line <> prettyPrec 0 t2
@@ -262,6 +267,22 @@ prettyTuple = tupled . map ppr . unnestTuple
   unnestTuple (TPair t1 t2) = t1 : unnestTuple t2
   unnestTuple t = [t]
 
+prettyDefinition :: Doc ann -> Var -> Maybe Polytype -> Term -> Doc ann
+prettyDefinition defName x mty t1 =
+  nest 2 . sep $
+    [ flatAlt
+        (defHead <> group defType <+> eqAndLambdaLine)
+        (defHead <> group defType' <+> defEqLambdas)
+    , ppr defBody
+    ]
+ where
+  (defBody, defLambdaList) = unchainLambdas t1
+  defHead = defName <+> pretty x
+  defType = maybe "" (\ty -> ":" <+> flatAlt (line <> indent 2 (ppr ty)) (ppr ty)) mty
+  defType' = maybe "" (\ty -> ":" <+> ppr ty) mty
+  defEqLambdas = hsep ("=" : map prettyLambda defLambdaList)
+  eqAndLambdaLine = if null defLambdaList then "=" else line <> defEqLambdas
+
 prettyPrecApp :: Int -> Term -> Term -> Doc a
 prettyPrecApp p t1 t2 =
   pparens (p > 10) $
@@ -413,7 +434,7 @@ fieldMismatchMsg expFs actFs =
 instance PrettyPrec InvalidAtomicReason where
   prettyPrec _ (TooManyTicks n) = "block could take too many ticks (" <> pretty n <> ")"
   prettyPrec _ AtomicDupingThing = "def, let, and lambda are not allowed"
-  prettyPrec _ (NonSimpleVarType _ ty) = "reference to variable with non-simple type" <+> ppr ty
+  prettyPrec _ (NonSimpleVarType _ ty) = "reference to variable with non-simple type" <+> ppr (prettyTextLine ty)
   prettyPrec _ NestedAtomic = "nested atomic block"
   prettyPrec _ LongConst = "commands that can take multiple ticks to execute are not allowed"
 

src/Swarm/Util.hs

@@ -27,6 +27,7 @@ module Swarm.Util (
   surfaceEmpty,
   applyWhen,
   hoistMaybe,
+  unsnocNE,
 
   -- * Directory utilities
   readFileMay,
@@ -239,6 +240,23 @@ applyWhen False _ x = x
 hoistMaybe :: (Applicative m) => Maybe b -> MaybeT m b
 hoistMaybe = MaybeT . pure
 
+-- | Like 'unsnoc', but for 'NonEmpty' so without the 'Maybe'
+--
+-- Taken from Cabal-syntax Distribution.Utils.Generic.
+--
+-- Example:
+-- >>> import Data.List.NonEmpty (NonEmpty ((:|)))
+-- >>> unsnocNE (1 :| [2, 3])
+-- ([1,2],3)
+--
+-- >>> unsnocNE (1 :| [])
+-- ([],1)
+unsnocNE :: NonEmpty a -> ([a], a)
+unsnocNE (x :| xs) = go x xs
+ where
+  go y [] = ([], y)
+  go y (z : zs) = let ~(ws, w) = go z zs in (y : ws, w)
+
 ------------------------------------------------------------
 -- Directory stuff
 

test/unit/TestLanguagePipeline.hs

@@ -11,15 +11,16 @@ import Control.Arrow ((&&&))
 import Control.Lens (toListOf)
 import Control.Lens.Plated (universe)
 import Data.Aeson (eitherDecode, encode)
-import Data.Either
 import Data.Maybe
 import Data.Text (Text)
 import Data.Text qualified as T
 import Data.Text.Encoding qualified as T
 import Swarm.Language.Module (Module (..))
+import Swarm.Language.Parse (readTerm)
 import Swarm.Language.Parse.QQ (tyQ)
 import Swarm.Language.Pipeline (ProcessedTerm (..), processTerm)
 import Swarm.Language.Pipeline.QQ (tmQ)
+import Swarm.Language.Pretty (prettyText)
 import Swarm.Language.Syntax
 import Swarm.Language.Typecheck (isSimpleUType)
 import Swarm.Language.Types
@@ -161,8 +162,8 @@ testLanguagePipeline =
         ]
     , testGroup
         "json encoding"
-        [ testCase "simple expr" (roundTrip "42 + 43")
-        , testCase "module def" (roundTrip "def x = 41 end; def y = 42 end")
+        [ testCase "simple expr" (roundTripTerm "42 + 43")
+        , testCase "module def" (roundTripTerm "def x = 41 end;\ndef y = 42 end")
         ]
     , testGroup
         "atomic - #479"
@@ -397,13 +398,6 @@ testLanguagePipeline =
  where
   valid = flip process ""
 
-  roundTrip txt = assertEqual "roundtrip" term (decodeThrow $ encode term)
-   where
-    decodeThrow v = case eitherDecode v of
-      Left e -> error $ "Decoding of " <> from (T.decodeUtf8 (from v)) <> " failed with: " <> from e
-      Right x -> x
-    term = fromMaybe (error "") $ fromRight (error "") $ processTerm txt
-
   process :: Text -> Text -> Assertion
   process code expect = case processTerm code of
     Left e
@@ -417,3 +411,14 @@ testLanguagePipeline =
 
   getSyntax :: ProcessedTerm -> Syntax' Polytype
   getSyntax (ProcessedTerm (Module s _) _ _) = s
+
+-- | Check round tripping of term from and to text, then test ToJSON/FromJSON.
+roundTripTerm :: Text -> Assertion
+roundTripTerm txt = do
+  assertEqual "roundtrip (readTerm -> prettyText)" txt (prettyText term)
+  assertEqual "roundtrip (ToJSON -> FromJSON)" term (decodeThrow $ encode term)
+ where
+  decodeThrow v = case eitherDecode v of
+    Left e -> error $ "Decoding of " <> from (T.decodeUtf8 (from v)) <> " failed with: " <> from e
+    Right x -> x
+  term = fromMaybe (error "empty document") $ either (error . T.unpack) id $ readTerm txt