{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE PatternGuards       #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE ViewPatterns        #-}
-- |
-- Module      : Data.Array.Accelerate.Pretty.Print
-- Copyright   : [2008..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.Pretty.Print (

  PrettyAcc, ExtractAcc,
  prettyPreOpenAcc,
  prettyPreOpenAfun,
  prettyOpenExp, prettyExp,
  prettyOpenFun, prettyFun,
  prettyArray,
  prettyConst,
  prettyELhs,
  prettyALhs,

  -- ** Internals
  Adoc,
  Val(..),
  PrettyEnv(..),
  Context(..),
  Keyword(..),
  Operator(..),
  parensIf, needsParens,
  ansiKeyword,
  shiftwidth,
  context0,
  app,
  manifest, delayed,
  primOperator,
  isInfix,
  prj, sizeEnv,

) where

import Data.Array.Accelerate.AST                                    hiding ( Direction )
import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.AST.LeftHandSide
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Elt
import Data.Array.Accelerate.Representation.Stencil
import Data.Array.Accelerate.Representation.Tag
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Sugar.Foreign
import Data.Array.Accelerate.Type
import qualified Data.Array.Accelerate.AST                          as AST

import Data.Char
import Data.String
import Data.Text.Prettyprint.Doc
import Data.Text.Prettyprint.Doc.Render.Terminal
import Prelude                                                      hiding ( exp )


-- Implementation
-- --------------

type PrettyAcc  acc = forall aenv a. Context -> Val aenv -> acc aenv a -> Adoc
type ExtractAcc acc = forall aenv a. acc aenv a -> PreOpenAcc acc aenv a

type Adoc = Doc Keyword

data Keyword
  = Statement     -- do | case of | let in
  | Conditional   -- if then else
  | Manifest      -- collective operations (kernel functions)
  | Delayed       -- fused operators
  deriving (Keyword -> Keyword -> Bool
(Keyword -> Keyword -> Bool)
-> (Keyword -> Keyword -> Bool) -> Eq Keyword
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Keyword -> Keyword -> Bool
$c/= :: Keyword -> Keyword -> Bool
== :: Keyword -> Keyword -> Bool
$c== :: Keyword -> Keyword -> Bool
Eq, Int -> Keyword -> ShowS
[Keyword] -> ShowS
Keyword -> String
(Int -> Keyword -> ShowS)
-> (Keyword -> String) -> ([Keyword] -> ShowS) -> Show Keyword
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Keyword] -> ShowS
$cshowList :: [Keyword] -> ShowS
show :: Keyword -> String
$cshow :: Keyword -> String
showsPrec :: Int -> Keyword -> ShowS
$cshowsPrec :: Int -> Keyword -> ShowS
Show)

let_, in_ :: Adoc
let_ :: Adoc
let_ = Keyword -> Adoc -> Adoc
forall ann. ann -> Doc ann -> Doc ann
annotate Keyword
Statement Adoc
"let"
in_ :: Adoc
in_  = Keyword -> Adoc -> Adoc
forall ann. ann -> Doc ann -> Doc ann
annotate Keyword
Statement Adoc
"in"

case_, of_ :: Adoc
case_ :: Adoc
case_ = Keyword -> Adoc -> Adoc
forall ann. ann -> Doc ann -> Doc ann
annotate Keyword
Statement Adoc
"case"
of_ :: Adoc
of_   = Keyword -> Adoc -> Adoc
forall ann. ann -> Doc ann -> Doc ann
annotate Keyword
Statement Adoc
"of"

if_, then_, else_ :: Adoc
if_ :: Adoc
if_   = Keyword -> Adoc -> Adoc
forall ann. ann -> Doc ann -> Doc ann
annotate Keyword
Statement Adoc
"if"
then_ :: Adoc
then_ = Keyword -> Adoc -> Adoc
forall ann. ann -> Doc ann -> Doc ann
annotate Keyword
Statement Adoc
"then"
else_ :: Adoc
else_ = Keyword -> Adoc -> Adoc
forall ann. ann -> Doc ann -> Doc ann
annotate Keyword
Statement Adoc
"else"

manifest :: Operator -> Adoc
manifest :: Operator -> Adoc
manifest = Keyword -> Adoc -> Adoc
forall ann. ann -> Doc ann -> Doc ann
annotate Keyword
Manifest (Adoc -> Adoc) -> (Operator -> Adoc) -> Operator -> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Operator -> Adoc
opName

delayed :: Operator -> Adoc
delayed :: Operator -> Adoc
delayed = Keyword -> Adoc -> Adoc
forall ann. ann -> Doc ann -> Doc ann
annotate Keyword
Delayed (Adoc -> Adoc) -> (Operator -> Adoc) -> Operator -> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Operator -> Adoc
opName

ansiKeyword :: Keyword -> AnsiStyle
ansiKeyword :: Keyword -> AnsiStyle
ansiKeyword Keyword
Statement   = Color -> AnsiStyle
colorDull Color
Yellow
ansiKeyword Keyword
Conditional = Color -> AnsiStyle
colorDull Color
Yellow
ansiKeyword Keyword
Manifest    = Color -> AnsiStyle
color Color
Blue
ansiKeyword Keyword
Delayed     = Color -> AnsiStyle
color Color
Green


-- Array computations
-- ------------------

prettyPreOpenAfun
    :: forall acc aenv f.
       PrettyAcc acc
    -> Val aenv
    -> PreOpenAfun acc aenv f
    -> Adoc
prettyPreOpenAfun :: PrettyAcc acc -> Val aenv -> PreOpenAfun acc aenv f -> Adoc
prettyPreOpenAfun PrettyAcc acc
prettyAcc Val aenv
aenv0 = Adoc -> Val aenv -> PreOpenAfun acc aenv f -> Adoc
forall aenv' f'.
Adoc -> Val aenv' -> PreOpenAfun acc aenv' f' -> Adoc
next (Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'\\') Val aenv
aenv0
  where
    next :: Adoc -> Val aenv' -> PreOpenAfun acc aenv' f' -> Adoc
    next :: Adoc -> Val aenv' -> PreOpenAfun acc aenv' f' -> Adoc
next Adoc
vs Val aenv'
aenv (Abody acc aenv' f'
body)   = Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [Adoc
vs Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
"->", Context -> Val aenv' -> acc aenv' f' -> Adoc
PrettyAcc acc
prettyAcc Context
context0 Val aenv'
aenv acc aenv' f'
body])
    next Adoc
vs Val aenv'
aenv (Alam ALeftHandSide a aenv' aenv'
lhs PreOpenAfun acc aenv' t
lam) =
      let (Val aenv'
aenv', Adoc
lhs') = Bool
-> Val aenv' -> ALeftHandSide a aenv' aenv' -> (Val aenv', Adoc)
forall env (s :: * -> *) arrs env'.
Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyALhs Bool
True Val aenv'
aenv ALeftHandSide a aenv' aenv'
lhs
      in  Adoc -> Val aenv' -> PreOpenAfun acc aenv' t -> Adoc
forall aenv' f'.
Adoc -> Val aenv' -> PreOpenAfun acc aenv' f' -> Adoc
next (Adoc
vs Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
lhs' Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
forall ann. Doc ann
space) Val aenv'
aenv' PreOpenAfun acc aenv' t
lam

prettyPreOpenAcc
    :: forall acc aenv arrs.
       Context
    -> PrettyAcc acc
    -> ExtractAcc acc
    -> Val aenv
    -> PreOpenAcc acc aenv arrs
    -> Adoc
prettyPreOpenAcc :: Context
-> PrettyAcc acc
-> ExtractAcc acc
-> Val aenv
-> PreOpenAcc acc aenv arrs
-> Adoc
prettyPreOpenAcc Context
ctx PrettyAcc acc
prettyAcc ExtractAcc acc
extractAcc Val aenv
aenv PreOpenAcc acc aenv arrs
pacc =
  case PreOpenAcc acc aenv arrs
pacc of
    Avar (Var ArrayR (Array sh e)
_ Idx aenv (Array sh e)
idx)        -> Idx aenv (Array sh e) -> Val aenv -> Adoc
forall env t. Idx env t -> Val env -> Adoc
prj Idx aenv (Array sh e)
idx Val aenv
aenv
    Alet{}                  -> Context
-> PrettyAcc acc
-> ExtractAcc acc
-> Val aenv
-> PreOpenAcc acc aenv arrs
-> Adoc
forall (acc :: * -> * -> *) aenv arrs.
Context
-> PrettyAcc acc
-> ExtractAcc acc
-> Val aenv
-> PreOpenAcc acc aenv arrs
-> Adoc
prettyAlet Context
ctx PrettyAcc acc
prettyAcc ExtractAcc acc
extractAcc Val aenv
aenv PreOpenAcc acc aenv arrs
pacc
    Apair{}                 -> PrettyAcc acc
-> ExtractAcc acc -> Val aenv -> PreOpenAcc acc aenv arrs -> Adoc
forall (acc :: * -> * -> *) aenv arrs.
PrettyAcc acc
-> ExtractAcc acc -> Val aenv -> PreOpenAcc acc aenv arrs -> Adoc
prettyAtuple PrettyAcc acc
prettyAcc ExtractAcc acc
extractAcc Val aenv
aenv PreOpenAcc acc aenv arrs
pacc
    PreOpenAcc acc aenv arrs
Anil                    -> Adoc
"()"
    Apply ArraysR arrs
_ PreOpenAfun acc aenv (arrs1 -> arrs)
f acc aenv arrs1
a             -> Adoc
apply
      where
        op :: Operator
op    = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
">->" Fixity
Infix Associativity
L Int
1
        apply :: Adoc
apply = [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ PreOpenAfun acc aenv (arrs1 -> arrs) -> Adoc
forall f. PreOpenAfun acc aenv f -> Adoc
ppAF PreOpenAfun acc aenv (arrs1 -> arrs)
f, Adoc -> Adoc
forall ann. Doc ann -> Doc ann
group ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [Operator -> Adoc
opName Operator
op, acc aenv arrs1 -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv arrs1
a]) ]

    Acond Exp aenv PrimBool
p acc aenv arrs
t acc aenv arrs
e             -> Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
flatAlt Adoc
multi Adoc
single
      where
        p' :: Adoc
p' = Exp aenv PrimBool -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv PrimBool
p
        t' :: Adoc
t' = acc aenv arrs -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv arrs
t
        e' :: Adoc
e' = acc aenv arrs -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv arrs
e
        --
        single :: Adoc
single = Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Operator -> Bool
needsParens Context
ctx (Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"?|:" Fixity
Infix Associativity
N Int
0))
               (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Adoc
p', Adoc
"?|", Adoc
t', Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
':', Adoc
e' ]
        multi :: Adoc
multi  = Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
3
               (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vsep [ Adoc
if_ Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
p'
                      , Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Adoc
then_, Adoc
t' ])
                      , Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Adoc
else_, Adoc
e' ]) ]

    Aforeign ArraysR arrs
_ asm (as -> arrs)
ff PreAfun acc (as -> arrs)
_ acc aenv as
a        -> Operator
"aforeign"       Operator -> [Adoc] -> Adoc
.$ [ String -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty (asm (as -> arrs) -> String
forall (asm :: * -> *) args. Foreign asm => asm args -> String
strForeign asm (as -> arrs)
ff), acc aenv as -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv as
a ]
    Awhile PreOpenAfun acc aenv (arrs -> Scalar PrimBool)
p PreOpenAfun acc aenv (arrs -> arrs)
f acc aenv arrs
a             -> Operator
"awhile"         Operator -> [Adoc] -> Adoc
.$ [ PreOpenAfun acc aenv (arrs -> Scalar PrimBool) -> Adoc
forall f. PreOpenAfun acc aenv f -> Adoc
ppAF PreOpenAfun acc aenv (arrs -> Scalar PrimBool)
p, PreOpenAfun acc aenv (arrs -> arrs) -> Adoc
forall f. PreOpenAfun acc aenv f -> Adoc
ppAF PreOpenAfun acc aenv (arrs -> arrs)
f, acc aenv arrs -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv arrs
a ]
    Use ArrayR (Array sh e)
repr Array sh e
arr             -> Operator
"use"            Operator -> [Adoc] -> Adoc
.$ [ ArrayR (Array sh e) -> Array sh e -> Adoc
forall sh e. ArrayR (Array sh e) -> Array sh e -> Adoc
prettyArray ArrayR (Array sh e)
repr Array sh e
arr ]
    Unit TypeR e
_ Exp aenv e
e                 -> Operator
"unit"           Operator -> [Adoc] -> Adoc
.$ [ Exp aenv e -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv e
e ]
    Reshape ShapeR sh
_ Exp aenv sh
sh acc aenv (Array sh' e)
a           -> Operator
"reshape"        Operator -> [Adoc] -> Adoc
.$ [ Exp aenv sh -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv sh
sh, acc aenv (Array sh' e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh' e)
a ]
    Generate ArrayR (Array sh e)
_ Exp aenv sh
sh Fun aenv (sh -> e)
f          -> Operator
"generate"       Operator -> [Adoc] -> Adoc
.$ [ Exp aenv sh -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv sh
sh, Fun aenv (sh -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (sh -> e)
f ]
    Transform ArrayR (Array sh' b)
_ Exp aenv sh'
sh Fun aenv (sh' -> sh)
p Fun aenv (a -> b)
f acc aenv (Array sh a)
a     -> Operator
"transform"      Operator -> [Adoc] -> Adoc
.$ [ Exp aenv sh' -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv sh'
sh, Fun aenv (sh' -> sh) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (sh' -> sh)
p, Fun aenv (a -> b) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (a -> b)
f, acc aenv (Array sh a) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh a)
a ]
    Replicate SliceIndex slix sl co sh
_ Exp aenv slix
ix acc aenv (Array sl e)
a         -> Operator
"replicate"      Operator -> [Adoc] -> Adoc
.$ [ Exp aenv slix -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv slix
ix, acc aenv (Array sl e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sl e)
a ]
    Slice SliceIndex slix sl co sh
_ acc aenv (Array sh e)
a Exp aenv slix
ix             -> Operator
"slice"          Operator -> [Adoc] -> Adoc
.$ [ Exp aenv slix -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv slix
ix, acc aenv (Array sh e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh e)
a ]
    Map TypeR e'
_ Fun aenv (e -> e')
f acc aenv (Array sh e)
a                -> Operator
"map"            Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e -> e') -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e -> e')
f,  acc aenv (Array sh e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh e)
a ]
    ZipWith TypeR e3
_ Fun aenv (e1 -> e2 -> e3)
f acc aenv (Array sh e1)
a acc aenv (Array sh e2)
b          -> Operator
"zipWith"        Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e1 -> e2 -> e3) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e1 -> e2 -> e3)
f,  acc aenv (Array sh e1) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh e1)
a, acc aenv (Array sh e2) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh e2)
b ]
    Fold Fun aenv (e -> e -> e)
f (Just Exp aenv e
z) acc aenv (Array (sh, Int) e)
a        -> Operator
"fold"           Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e -> e -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e -> e -> e)
f,  Exp aenv e -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv e
z, acc aenv (Array (sh, Int) e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array (sh, Int) e)
a ]
    Fold Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
Nothing  acc aenv (Array (sh, Int) e)
a        -> Operator
"fold1"          Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e -> e -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e -> e -> e)
f,  acc aenv (Array (sh, Int) e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array (sh, Int) e)
a ]
    FoldSeg IntegralType i
_ Fun aenv (e -> e -> e)
f (Just Exp aenv e
z) acc aenv (Array (sh, Int) e)
a acc aenv (Segments i)
s -> Operator
"foldSeg"        Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e -> e -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e -> e -> e)
f,  Exp aenv e -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv e
z, acc aenv (Array (sh, Int) e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array (sh, Int) e)
a, acc aenv (Segments i) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Segments i)
s ]
    FoldSeg IntegralType i
_ Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
Nothing  acc aenv (Array (sh, Int) e)
a acc aenv (Segments i)
s -> Operator
"fold1Seg"       Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e -> e -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e -> e -> e)
f,  acc aenv (Array (sh, Int) e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array (sh, Int) e)
a, acc aenv (Segments i) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Segments i)
s ]
    Scan Direction
d Fun aenv (e -> e -> e)
f (Just Exp aenv e
z) acc aenv (Array (sh, Int) e)
a      -> String -> Direction -> String -> Operator
ppD String
"scan" Direction
d String
""  Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e -> e -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e -> e -> e)
f,  Exp aenv e -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv e
z, acc aenv (Array (sh, Int) e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array (sh, Int) e)
a ]
    Scan Direction
d Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
Nothing  acc aenv (Array (sh, Int) e)
a      -> String -> Direction -> String -> Operator
ppD String
"scan" Direction
d String
"1" Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e -> e -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e -> e -> e)
f,  acc aenv (Array (sh, Int) e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array (sh, Int) e)
a ]
    Scan' Direction
d Fun aenv (e -> e -> e)
f Exp aenv e
z acc aenv (Array (sh, Int) e)
a            -> String -> Direction -> String -> Operator
ppD String
"scan" Direction
d String
"'" Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e -> e -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e -> e -> e)
f,  Exp aenv e -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv e
z, acc aenv (Array (sh, Int) e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array (sh, Int) e)
a ]
    Permute Fun aenv (e -> e -> e)
f acc aenv (Array sh' e)
d Fun aenv (sh -> PrimMaybe sh')
p acc aenv (Array sh e)
s          -> Operator
"permute"        Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (e -> e -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (e -> e -> e)
f,  acc aenv (Array sh' e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh' e)
d, Fun aenv (sh -> PrimMaybe sh') -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (sh -> PrimMaybe sh')
p, acc aenv (Array sh e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh e)
s ]
    Backpermute ShapeR sh'
_ Exp aenv sh'
sh Fun aenv (sh' -> sh)
f acc aenv (Array sh e)
a     -> Operator
"backpermute"    Operator -> [Adoc] -> Adoc
.$ [ Exp aenv sh' -> Adoc
forall t. Exp aenv t -> Adoc
ppE Exp aenv sh'
sh, Fun aenv (sh' -> sh) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (sh' -> sh)
f, acc aenv (Array sh e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh e)
a ]
    Stencil StencilR sh e stencil
s TypeR e'
_ Fun aenv (stencil -> e')
f Boundary aenv (Array sh e)
b acc aenv (Array sh e)
a        -> Operator
"stencil"        Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (stencil -> e') -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (stencil -> e')
f,  TypeR e -> Boundary aenv (Array sh e) -> Adoc
forall sh e. TypeR e -> Boundary aenv (Array sh e) -> Adoc
ppB (StencilR sh e stencil -> TypeR e
forall sh e pat. StencilR sh e pat -> TypeR e
stencilEltR StencilR sh e stencil
s) Boundary aenv (Array sh e)
b, acc aenv (Array sh e) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh e)
a ]
    Stencil2 StencilR sh a stencil1
s1 StencilR sh b stencil2
s2 TypeR c
_ Fun aenv (stencil1 -> stencil2 -> c)
f Boundary aenv (Array sh a)
b1 acc aenv (Array sh a)
a1 Boundary aenv (Array sh b)
b2 acc aenv (Array sh b)
a2
                             -> Operator
"stencil2"       Operator -> [Adoc] -> Adoc
.$ [ Fun aenv (stencil1 -> stencil2 -> c) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (stencil1 -> stencil2 -> c)
f,  TypeR a -> Boundary aenv (Array sh a) -> Adoc
forall sh e. TypeR e -> Boundary aenv (Array sh e) -> Adoc
ppB (StencilR sh a stencil1 -> TypeR a
forall sh e pat. StencilR sh e pat -> TypeR e
stencilEltR StencilR sh a stencil1
s1) Boundary aenv (Array sh a)
b1, acc aenv (Array sh a) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh a)
a1, TypeR b -> Boundary aenv (Array sh b) -> Adoc
forall sh e. TypeR e -> Boundary aenv (Array sh e) -> Adoc
ppB (StencilR sh b stencil2 -> TypeR b
forall sh e pat. StencilR sh e pat -> TypeR e
stencilEltR StencilR sh b stencil2
s2) Boundary aenv (Array sh b)
b2, acc aenv (Array sh b) -> Adoc
forall a. acc aenv a -> Adoc
ppA acc aenv (Array sh b)
a2 ]
  where
    infixr 0 .$
    Operator
f .$ :: Operator -> [Adoc] -> Adoc
.$ [Adoc]
xs
      = Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Operator -> Bool
needsParens Context
ctx Operator
f)
      (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep (Operator -> Adoc
manifest Operator
f Adoc -> [Adoc] -> [Adoc]
forall a. a -> [a] -> [a]
: [Adoc]
xs))

    ppA :: acc aenv a -> Adoc
    ppA :: acc aenv a -> Adoc
ppA = Context -> Val aenv -> acc aenv a -> Adoc
PrettyAcc acc
prettyAcc Context
app Val aenv
aenv

    ppAF :: PreOpenAfun acc aenv f -> Adoc
    ppAF :: PreOpenAfun acc aenv f -> Adoc
ppAF = Adoc -> Adoc
forall ann. Doc ann -> Doc ann
parens (Adoc -> Adoc)
-> (PreOpenAfun acc aenv f -> Adoc)
-> PreOpenAfun acc aenv f
-> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrettyAcc acc -> Val aenv -> PreOpenAfun acc aenv f -> Adoc
forall (acc :: * -> * -> *) aenv f.
PrettyAcc acc -> Val aenv -> PreOpenAfun acc aenv f -> Adoc
prettyPreOpenAfun PrettyAcc acc
prettyAcc Val aenv
aenv

    ppE :: Exp aenv t -> Adoc
    ppE :: Exp aenv t -> Adoc
ppE = Context -> Val () -> Val aenv -> Exp aenv t -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
app Val ()
Empty Val aenv
aenv

    ppF :: Fun aenv t -> Adoc
    ppF :: Fun aenv t -> Adoc
ppF = Adoc -> Adoc
forall ann. Doc ann -> Doc ann
parens (Adoc -> Adoc) -> (Fun aenv t -> Adoc) -> Fun aenv t -> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Val () -> Val aenv -> Fun aenv t -> Adoc
forall env aenv f.
Val env -> Val aenv -> OpenFun env aenv f -> Adoc
prettyOpenFun Val ()
Empty Val aenv
aenv

    ppB :: forall sh e.
           TypeR e
        -> Boundary aenv (Array sh e)
        -> Adoc
    ppB :: TypeR e -> Boundary aenv (Array sh e) -> Adoc
ppB TypeR e
_  Boundary aenv (Array sh e)
Clamp        = Adoc
"clamp"
    ppB TypeR e
_  Boundary aenv (Array sh e)
Mirror       = Adoc
"mirror"
    ppB TypeR e
_  Boundary aenv (Array sh e)
Wrap         = Adoc
"wrap"
    ppB TypeR e
tp (Constant e
e) = TypeR e -> e -> Adoc
forall e. TypeR e -> e -> Adoc
prettyConst TypeR e
tp e
e
e
    ppB TypeR e
_  (Function Fun aenv (sh -> e)
f) = Fun aenv (sh -> e) -> Adoc
forall t. Fun aenv t -> Adoc
ppF Fun aenv (sh -> e)
f

    ppD :: String -> AST.Direction -> String -> Operator
    ppD :: String -> Direction -> String -> Operator
ppD String
f Direction
AST.LeftToRight String
k = String -> Operator
forall a. IsString a => String -> a
fromString (String
f String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"l" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
k)
    ppD String
f Direction
AST.RightToLeft String
k = String -> Operator
forall a. IsString a => String -> a
fromString (String
f String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"r" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
k)


prettyAlet
    :: forall acc aenv arrs.
       Context
    -> PrettyAcc acc
    -> ExtractAcc acc
    -> Val aenv
    -> PreOpenAcc acc aenv arrs
    -> Adoc
prettyAlet :: Context
-> PrettyAcc acc
-> ExtractAcc acc
-> Val aenv
-> PreOpenAcc acc aenv arrs
-> Adoc
prettyAlet Context
ctx PrettyAcc acc
prettyAcc ExtractAcc acc
extractAcc Val aenv
aenv0
  = Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Operator -> Bool
needsParens Context
ctx Operator
"let")
  (Adoc -> Adoc)
-> (PreOpenAcc acc aenv arrs -> Adoc)
-> PreOpenAcc acc aenv arrs
-> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Adoc -> Adoc
forall ann. Doc ann -> Doc ann
align (Adoc -> Adoc)
-> (PreOpenAcc acc aenv arrs -> Adoc)
-> PreOpenAcc acc aenv arrs
-> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Adoc], Adoc) -> Adoc
wrap (([Adoc], Adoc) -> Adoc)
-> (PreOpenAcc acc aenv arrs -> ([Adoc], Adoc))
-> PreOpenAcc acc aenv arrs
-> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Val aenv -> PreOpenAcc acc aenv arrs -> ([Adoc], Adoc)
forall aenv' a.
Val aenv' -> PreOpenAcc acc aenv' a -> ([Adoc], Adoc)
collect Val aenv
aenv0
  where
    collect :: Val aenv' -> PreOpenAcc acc aenv' a -> ([Adoc], Adoc)
    collect :: Val aenv' -> PreOpenAcc acc aenv' a -> ([Adoc], Adoc)
collect Val aenv'
aenv =
      \case
        Alet ALeftHandSide bndArrs aenv' aenv'
lhs acc aenv' bndArrs
a1 acc aenv' a
a2 ->
          let (Val aenv'
aenv', Adoc
v)      = Bool
-> Val aenv'
-> ALeftHandSide bndArrs aenv' aenv'
-> (Val aenv', Adoc)
forall env (s :: * -> *) arrs env'.
Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyALhs Bool
False Val aenv'
aenv ALeftHandSide bndArrs aenv' aenv'
lhs
              a1' :: Adoc
a1'             = Val aenv' -> acc aenv' bndArrs -> Adoc
forall aenv' a. Val aenv' -> acc aenv' a -> Adoc
ppA Val aenv'
aenv acc aenv' bndArrs
a1
              bnd :: Adoc
bnd | acc aenv' bndArrs -> Bool
forall aenv' a. acc aenv' a -> Bool
isAlet acc aenv' bndArrs
a1 = Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
nest Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vsep [Adoc
v Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
forall ann. Doc ann
equals, Adoc
a1'])
                  | Bool
otherwise = Adoc
v Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc -> Adoc
forall ann. Doc ann -> Doc ann
align (Adoc
forall ann. Doc ann
equals Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
a1')
              ([Adoc]
bnds, Adoc
body)    = Val aenv' -> PreOpenAcc acc aenv' a -> ([Adoc], Adoc)
forall aenv' a.
Val aenv' -> PreOpenAcc acc aenv' a -> ([Adoc], Adoc)
collect Val aenv'
aenv' (acc aenv' a -> PreOpenAcc acc aenv' a
ExtractAcc acc
extractAcc acc aenv' a
a2)
          in
          (Adoc
bndAdoc -> [Adoc] -> [Adoc]
forall a. a -> [a] -> [a]
:[Adoc]
bnds, Adoc
body)
        --
        PreOpenAcc acc aenv' a
next       -> ([], Context
-> PrettyAcc acc
-> ExtractAcc acc
-> Val aenv'
-> PreOpenAcc acc aenv' a
-> Adoc
forall (acc :: * -> * -> *) aenv arrs.
Context
-> PrettyAcc acc
-> ExtractAcc acc
-> Val aenv
-> PreOpenAcc acc aenv arrs
-> Adoc
prettyPreOpenAcc Context
context0 PrettyAcc acc
prettyAcc ExtractAcc acc
extractAcc Val aenv'
aenv PreOpenAcc acc aenv' a
next)

    isAlet :: acc aenv' a -> Bool
    isAlet :: acc aenv' a -> Bool
isAlet (acc aenv' a -> PreOpenAcc acc aenv' a
ExtractAcc acc
extractAcc -> Alet{}) = Bool
True
    isAlet acc aenv' a
_                      = Bool
False

    ppA :: Val aenv' -> acc aenv' a -> Adoc
    ppA :: Val aenv' -> acc aenv' a -> Adoc
ppA = Context -> Val aenv' -> acc aenv' a -> Adoc
PrettyAcc acc
prettyAcc Context
context0

    wrap :: ([Adoc], Adoc) -> Adoc
    wrap :: ([Adoc], Adoc) -> Adoc
wrap ([],   Adoc
body) = Adoc
body  -- shouldn't happen!
    wrap ([Adoc
b],  Adoc
body)
      = [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
nest Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [Adoc
let_, Adoc
b]), Adoc
in_, Adoc
body ]
    wrap ([Adoc]
bnds, Adoc
body)
      = [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vsep [ Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
nest Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vsep (Adoc
let_Adoc -> [Adoc] -> [Adoc]
forall a. a -> [a] -> [a]
:[Adoc]
bnds))
             , Adoc
in_
             , Adoc
body
             ]

prettyAtuple
    :: forall acc aenv arrs.
       PrettyAcc acc
    -> ExtractAcc acc
    -> Val aenv
    -> PreOpenAcc acc aenv arrs
    -> Adoc
prettyAtuple :: PrettyAcc acc
-> ExtractAcc acc -> Val aenv -> PreOpenAcc acc aenv arrs -> Adoc
prettyAtuple PrettyAcc acc
prettyAcc ExtractAcc acc
extractAcc Val aenv
aenv0 PreOpenAcc acc aenv arrs
acc = case PreOpenAcc acc aenv arrs -> Maybe [Adoc]
forall arrs'. PreOpenAcc acc aenv arrs' -> Maybe [Adoc]
collect PreOpenAcc acc aenv arrs
acc of
    Maybe [Adoc]
Nothing  -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann
align (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ PreOpenAcc acc aenv arrs -> Adoc
forall arrs'. PreOpenAcc acc aenv arrs' -> Adoc
ppPair PreOpenAcc acc aenv arrs
acc
    Just [Adoc]
tup ->
      case [Adoc]
tup of
        []  -> Adoc
"()"
        [Adoc
t] -> Adoc
t
        [Adoc]
_   -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann
align (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ Adoc
"T" Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Int -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty ([Adoc] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Adoc]
tup) Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [Adoc]
tup
  where
    ppPair :: PreOpenAcc acc aenv arrs' -> Adoc
    ppPair :: PreOpenAcc acc aenv arrs' -> Adoc
ppPair (Apair acc aenv as
a1 acc aenv bs
a2) = Adoc
"(" Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> PreOpenAcc acc aenv as -> Adoc
forall arrs'. PreOpenAcc acc aenv arrs' -> Adoc
ppPair (acc aenv as -> PreOpenAcc acc aenv as
ExtractAcc acc
extractAcc acc aenv as
a1) Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
"," Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Context -> Val aenv -> acc aenv bs -> Adoc
PrettyAcc acc
prettyAcc Context
context0 Val aenv
aenv0 acc aenv bs
a2 Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
")"
    ppPair PreOpenAcc acc aenv arrs'
a             = Context
-> PrettyAcc acc
-> ExtractAcc acc
-> Val aenv
-> PreOpenAcc acc aenv arrs'
-> Adoc
forall (acc :: * -> * -> *) aenv arrs.
Context
-> PrettyAcc acc
-> ExtractAcc acc
-> Val aenv
-> PreOpenAcc acc aenv arrs
-> Adoc
prettyPreOpenAcc Context
context0 PrettyAcc acc
prettyAcc ExtractAcc acc
extractAcc Val aenv
aenv0 PreOpenAcc acc aenv arrs'
a

    collect :: PreOpenAcc acc aenv arrs' -> Maybe [Adoc]
    collect :: PreOpenAcc acc aenv arrs' -> Maybe [Adoc]
collect PreOpenAcc acc aenv arrs'
Anil          = [Adoc] -> Maybe [Adoc]
forall a. a -> Maybe a
Just []
    collect (Apair acc aenv as
a1 acc aenv bs
a2)
      | Just [Adoc]
tup <- PreOpenAcc acc aenv as -> Maybe [Adoc]
forall arrs'. PreOpenAcc acc aenv arrs' -> Maybe [Adoc]
collect (PreOpenAcc acc aenv as -> Maybe [Adoc])
-> PreOpenAcc acc aenv as -> Maybe [Adoc]
forall a b. (a -> b) -> a -> b
$ acc aenv as -> PreOpenAcc acc aenv as
ExtractAcc acc
extractAcc acc aenv as
a1
                          = [Adoc] -> Maybe [Adoc]
forall a. a -> Maybe a
Just ([Adoc] -> Maybe [Adoc]) -> [Adoc] -> Maybe [Adoc]
forall a b. (a -> b) -> a -> b
$ [Adoc]
tup [Adoc] -> [Adoc] -> [Adoc]
forall a. [a] -> [a] -> [a]
++ [Context -> Val aenv -> acc aenv bs -> Adoc
PrettyAcc acc
prettyAcc Context
app Val aenv
aenv0 acc aenv bs
a2]
    collect PreOpenAcc acc aenv arrs'
_             = Maybe [Adoc]
forall a. Maybe a
Nothing

-- TODO: Should we also print the types of the declared variables? And the types of wildcards?
prettyALhs :: Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyALhs :: Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyALhs Bool
requiresParens = Bool
-> Char
-> Val env
-> LeftHandSide s arrs env env'
-> (Val env', Adoc)
forall (s :: * -> *) env env' arrs.
Bool
-> Char
-> Val env
-> LeftHandSide s arrs env env'
-> (Val env', Adoc)
prettyLhs Bool
requiresParens Char
'a'

prettyELhs :: Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyELhs :: Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyELhs Bool
requiresParens = Bool
-> Char
-> Val env
-> LeftHandSide s arrs env env'
-> (Val env', Adoc)
forall (s :: * -> *) env env' arrs.
Bool
-> Char
-> Val env
-> LeftHandSide s arrs env env'
-> (Val env', Adoc)
prettyLhs Bool
requiresParens Char
'x'

prettyLhs :: forall s env env' arrs. Bool -> Char -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyLhs :: Bool
-> Char
-> Val env
-> LeftHandSide s arrs env env'
-> (Val env', Adoc)
prettyLhs Bool
requiresParens Char
x Val env
env0 LeftHandSide s arrs env env'
lhs = case LeftHandSide s arrs env env' -> Maybe (Val env', [Adoc])
forall arrs' env''.
LeftHandSide s arrs' env env'' -> Maybe (Val env'', [Adoc])
collect LeftHandSide s arrs env env'
lhs of
  Maybe (Val env', [Adoc])
Nothing          -> LeftHandSide s arrs env env' -> (Val env', Adoc)
forall arrs' env''.
LeftHandSide s arrs' env env'' -> (Val env'', Adoc)
ppPair LeftHandSide s arrs env env'
lhs
  Just (Val env'
env1, [Adoc]
tup) ->
    case [Adoc]
tup of
      []  -> (Val env'
env1, Adoc
"()")
      [Adoc]
_   -> (Val env'
env1, Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf Bool
requiresParens (Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'T' Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Int -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty ([Adoc] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Adoc]
tup) Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [Adoc]
tup))
  where
    ppPair :: LeftHandSide s arrs' env env'' -> (Val env'', Adoc)
    ppPair :: LeftHandSide s arrs' env env'' -> (Val env'', Adoc)
ppPair LeftHandSide s arrs' env env''
LeftHandSideUnit       = (Val env
Val env''
env0, Adoc
"()")
    ppPair LeftHandSideWildcard{} = (Val env
Val env''
env0, Adoc
"_")
    ppPair LeftHandSideSingle{}   = (Val env
env0 Val env -> Adoc -> Val (env, arrs')
forall env t. Val env -> Adoc -> Val (env, t)
`Push` Adoc
v, Adoc
v)
      where
        v :: Adoc
v = Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
x Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Int -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty (Val env -> Int
forall env. Val env -> Int
sizeEnv Val env
env0)
    ppPair (LeftHandSidePair LeftHandSide s v1 env env'
a LeftHandSide s v2 env' env''
b)          = (Val env''
env2, [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
tupled [Adoc
doc1, Adoc
doc2])
      where
        (Val env'
env1, Adoc
doc1) = LeftHandSide s v1 env env' -> (Val env', Adoc)
forall arrs' env''.
LeftHandSide s arrs' env env'' -> (Val env'', Adoc)
ppPair LeftHandSide s v1 env env'
a
        (Val env''
env2, Adoc
doc2) = Bool
-> Char
-> Val env'
-> LeftHandSide s v2 env' env''
-> (Val env'', Adoc)
forall (s :: * -> *) env env' arrs.
Bool
-> Char
-> Val env
-> LeftHandSide s arrs env env'
-> (Val env', Adoc)
prettyLhs Bool
False Char
x Val env'
env1 LeftHandSide s v2 env' env''
b

    collect :: LeftHandSide s arrs' env env'' -> Maybe (Val env'', [Adoc])
    collect :: LeftHandSide s arrs' env env'' -> Maybe (Val env'', [Adoc])
collect (LeftHandSidePair LeftHandSide s v1 env env'
l1 LeftHandSide s v2 env' env''
l2)
      | Just (Val env'
env1, [Adoc]
tup ) <- LeftHandSide s v1 env env' -> Maybe (Val env', [Adoc])
forall arrs' env''.
LeftHandSide s arrs' env env'' -> Maybe (Val env'', [Adoc])
collect LeftHandSide s v1 env env'
l1
      ,      (Val env''
env2, Adoc
doc2) <- Bool
-> Char
-> Val env'
-> LeftHandSide s v2 env' env''
-> (Val env'', Adoc)
forall (s :: * -> *) env env' arrs.
Bool
-> Char
-> Val env
-> LeftHandSide s arrs env env'
-> (Val env', Adoc)
prettyLhs Bool
True Char
x Val env'
env1 LeftHandSide s v2 env' env''
l2 = (Val env'', [Adoc]) -> Maybe (Val env'', [Adoc])
forall a. a -> Maybe a
Just (Val env''
env2, [Adoc]
tup [Adoc] -> [Adoc] -> [Adoc]
forall a. [a] -> [a] -> [a]
++ [Adoc
doc2])
    collect (LeftHandSideWildcard TupR s arrs'
TupRunit) = (Val env, [Adoc]) -> Maybe (Val env, [Adoc])
forall a. a -> Maybe a
Just (Val env
env0, [])
    collect LeftHandSide s arrs' env env''
_ = Maybe (Val env'', [Adoc])
forall a. Maybe a
Nothing

prettyArray :: ArrayR (Array sh e) -> Array sh e -> Adoc
prettyArray :: ArrayR (Array sh e) -> Array sh e -> Adoc
prettyArray aR :: ArrayR (Array sh e)
aR@(ArrayR ShapeR sh
_ TypeR e
eR) = Adoc -> Adoc
forall ann. Doc ann -> Doc ann
parens (Adoc -> Adoc) -> (Array sh e -> Adoc) -> Array sh e -> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Adoc
forall a. IsString a => String -> a
fromString (String -> Adoc) -> (Array sh e -> String) -> Array sh e -> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (e -> ShowS) -> ArrayR (Array sh e) -> Array sh e -> String
forall e sh.
(e -> ShowS) -> ArrayR (Array sh e) -> Array sh e -> String
showArray (TypeR e -> e -> ShowS
forall e. TypeR e -> e -> ShowS
showsElt TypeR e
eR) ArrayR (Array sh e)
ArrayR (Array sh e)
aR


-- Scalar expressions
-- ------------------

prettyFun :: Val aenv -> Fun aenv f -> Adoc
prettyFun :: Val aenv -> Fun aenv f -> Adoc
prettyFun = Val () -> Val aenv -> Fun aenv f -> Adoc
forall env aenv f.
Val env -> Val aenv -> OpenFun env aenv f -> Adoc
prettyOpenFun Val ()
Empty

prettyExp :: Val aenv -> Exp aenv t -> Adoc
prettyExp :: Val aenv -> Exp aenv t -> Adoc
prettyExp = Context -> Val () -> Val aenv -> Exp aenv t -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
context0 Val ()
Empty

prettyOpenFun
    :: forall env aenv f.
       Val env
    -> Val aenv
    -> OpenFun env aenv f
    -> Adoc
prettyOpenFun :: Val env -> Val aenv -> OpenFun env aenv f -> Adoc
prettyOpenFun Val env
env0 Val aenv
aenv = Adoc -> Val env -> OpenFun env aenv f -> Adoc
forall env' f'. Adoc -> Val env' -> OpenFun env' aenv f' -> Adoc
next (Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'\\') Val env
env0
  where
    next :: Adoc -> Val env' -> OpenFun env' aenv f' -> Adoc
    next :: Adoc -> Val env' -> OpenFun env' aenv f' -> Adoc
next Adoc
vs Val env'
env (Body OpenExp env' aenv f'
body)
      --   PrimApp f x                             <- body
      -- , op                                      <- primOperator f
      -- , isInfix op
      -- , Tuple (NilTup `SnocTup` a `SnocTup` b)  <- x
      -- , Var (SuccIdx ZeroIdx)                   <- a
      -- , Var ZeroIdx                             <- b
      -- = opName op -- surrounding context will add parens
      --
      = Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Adoc
vs Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
"->"
                             , Context -> Val env' -> Val aenv -> OpenExp env' aenv f' -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
context0 Val env'
env Val aenv
aenv OpenExp env' aenv f'
body])
    next Adoc
vs Val env'
env (Lam ELeftHandSide a env' env'
lhs OpenFun env' aenv t
lam) =
      let (Val env'
env', Adoc
lhs') = Bool -> Val env' -> ELeftHandSide a env' env' -> (Val env', Adoc)
forall env (s :: * -> *) arrs env'.
Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyELhs Bool
True Val env'
env ELeftHandSide a env' env'
lhs
      in  Adoc -> Val env' -> OpenFun env' aenv t -> Adoc
forall env' f'. Adoc -> Val env' -> OpenFun env' aenv f' -> Adoc
next (Adoc
vs Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
lhs' Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
forall ann. Doc ann
space) Val env'
env' OpenFun env' aenv t
lam

prettyOpenExp
    :: forall env aenv t.
       Context
    -> Val env
    -> Val aenv
    -> OpenExp env aenv t
    -> Adoc
prettyOpenExp :: Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
ctx Val env
env Val aenv
aenv OpenExp env aenv t
exp =
  case OpenExp env aenv t
exp of
    Evar (Var ScalarType t
_ Idx env t
idx)      -> Idx env t -> Val env -> Adoc
forall env t. Idx env t -> Val env -> Adoc
prj Idx env t
idx Val env
env
    Let{}                 -> Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyLet Context
ctx Val env
env Val aenv
aenv OpenExp env aenv t
exp
    PrimApp PrimFun (a -> t)
f OpenExp env aenv a
x
      | OpenExp env aenv t1
a `Pair` OpenExp env aenv t2
b <- OpenExp env aenv a
x   -> Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
ppF2 Operator
op  (OpenExp env aenv t1 -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv t1
a) (OpenExp env aenv t2 -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv t2
b)
      | Bool
otherwise         -> Operator -> (Context -> Adoc) -> Adoc
ppF1 Operator
op' (OpenExp env aenv a -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv a
x)
      where
        op :: Operator
op  = PrimFun (a -> t) -> Operator
forall a. PrimFun a -> Operator
primOperator PrimFun (a -> t)
f
        op' :: Operator
op' = Operator -> Bool
isInfix Operator
op Bool -> (Operator, Operator) -> Operator
forall a. Bool -> (a, a) -> a
? (Adoc -> Fixity -> Associativity -> Int -> Operator
Operator (Adoc -> Adoc
forall ann. Doc ann -> Doc ann
parens (Operator -> Adoc
opName Operator
op)) Fixity
App Associativity
L Int
10, Operator
op)
    --
    PrimConst PrimConst t
c           -> PrimConst t -> Adoc
forall a. PrimConst a -> Adoc
prettyPrimConst PrimConst t
c
    Const ScalarType t
tp t
c            -> TypeR t -> t -> Adoc
forall e. TypeR e -> e -> Adoc
prettyConst (ScalarType t -> TypeR t
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ScalarType t
tp) t
c
    Pair{}                -> Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyTuple Context
ctx Val env
env Val aenv
aenv OpenExp env aenv t
exp
    OpenExp env aenv t
Nil                   -> Adoc
"()"
    VecPack   VecR n s tup
_ OpenExp env aenv tup
e         -> Operator -> (Context -> Adoc) -> Adoc
ppF1 Operator
"pack"   (OpenExp env aenv tup -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv tup
e)
    VecUnpack VecR n s t
_ OpenExp env aenv (Vec n s)
e         -> Operator -> (Context -> Adoc) -> Adoc
ppF1 Operator
"unpack" (OpenExp env aenv (Vec n s) -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv (Vec n s)
e)
    Case OpenExp env aenv PrimBool
x [(PrimBool, OpenExp env aenv t)]
xs Maybe (OpenExp env aenv t)
d           -> Val env
-> Val aenv
-> OpenExp env aenv PrimBool
-> [(PrimBool, OpenExp env aenv t)]
-> Maybe (OpenExp env aenv t)
-> Adoc
forall env aenv a b.
Val env
-> Val aenv
-> OpenExp env aenv a
-> [(PrimBool, OpenExp env aenv b)]
-> Maybe (OpenExp env aenv b)
-> Adoc
prettyCase Val env
env Val aenv
aenv OpenExp env aenv PrimBool
x [(PrimBool, OpenExp env aenv t)]
xs Maybe (OpenExp env aenv t)
d
    Cond OpenExp env aenv PrimBool
p OpenExp env aenv t
t OpenExp env aenv t
e            -> Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
flatAlt Adoc
multi Adoc
single
      where
        p' :: Adoc
p' = OpenExp env aenv PrimBool -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv PrimBool
p Context
context0
        t' :: Adoc
t' = OpenExp env aenv t -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv t
t Context
context0
        e' :: Adoc
e' = OpenExp env aenv t -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv t
e Context
context0
        --
        single :: Adoc
single = Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Operator -> Bool
needsParens Context
ctx (Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"?:" Fixity
Infix Associativity
N Int
0))
               (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Adoc
p', Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'?', Adoc
t', Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
':', Adoc
e' ]
        multi :: Adoc
multi  = Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
3
               (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vsep [ Adoc
if_ Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
p'
                      , Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Adoc
then_, Adoc
t' ])
                      , Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Adoc
else_, Adoc
e' ]) ]
    --
    IndexSlice SliceIndex slix t co sh
_ OpenExp env aenv slix
slix OpenExp env aenv sh
sh  -> Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
ppF2 Operator
"indexSlice"  (OpenExp env aenv slix -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv slix
slix) (OpenExp env aenv sh -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv sh
sh)
    IndexFull SliceIndex slix sl co t
_ OpenExp env aenv slix
slix OpenExp env aenv sl
sl   -> Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
ppF2 Operator
"indexFull"   (OpenExp env aenv slix -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv slix
slix) (OpenExp env aenv sl -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv sl
sl)
    ToIndex ShapeR sh
_ OpenExp env aenv sh
sh OpenExp env aenv sh
ix       -> Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
ppF2 Operator
"toIndex"     (OpenExp env aenv sh -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv sh
sh) (OpenExp env aenv sh -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv sh
ix)
    FromIndex ShapeR t
_ OpenExp env aenv t
sh OpenExp env aenv Int
ix     -> Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
ppF2 Operator
"fromIndex"   (OpenExp env aenv t -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv t
sh) (OpenExp env aenv Int -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv Int
ix)
    While OpenFun env aenv (t -> PrimBool)
p OpenFun env aenv (t -> t)
f OpenExp env aenv t
x           -> Operator
-> (Context -> Adoc)
-> (Context -> Adoc)
-> (Context -> Adoc)
-> Adoc
ppF3 Operator
"while"       (OpenFun env aenv (t -> PrimBool) -> Context -> Adoc
forall f. OpenFun env aenv f -> Context -> Adoc
ppF OpenFun env aenv (t -> PrimBool)
p) (OpenFun env aenv (t -> t) -> Context -> Adoc
forall f. OpenFun env aenv f -> Context -> Adoc
ppF OpenFun env aenv (t -> t)
f) (OpenExp env aenv t -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv t
x)
    Foreign TypeR t
_ asm (x -> t)
ff Fun () (x -> t)
_ OpenExp env aenv x
e      -> Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
ppF2 Operator
"foreign"     (\Context
_ -> String -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty (asm (x -> t) -> String
forall (asm :: * -> *) args. Foreign asm => asm args -> String
strForeign asm (x -> t)
ff)) (OpenExp env aenv x -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv x
e)
    Shape ArrayVar aenv (Array t e)
arr             -> Operator -> (Context -> Adoc) -> Adoc
ppF1 Operator
"shape"       (ArrayVar aenv (Array t e) -> Context -> Adoc
forall a. ArrayVar aenv a -> Context -> Adoc
ppA ArrayVar aenv (Array t e)
arr)
    ShapeSize ShapeR dim
_ OpenExp env aenv dim
sh        -> Operator -> (Context -> Adoc) -> Adoc
ppF1 Operator
"shapeSize"   (OpenExp env aenv dim -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv dim
sh)
    Index ArrayVar aenv (Array dim t)
arr OpenExp env aenv dim
ix          -> Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
ppF2 (Adoc -> Fixity -> Associativity -> Int -> Operator
Operator (Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'!') Fixity
Infix Associativity
L Int
9) (ArrayVar aenv (Array dim t) -> Context -> Adoc
forall a. ArrayVar aenv a -> Context -> Adoc
ppA ArrayVar aenv (Array dim t)
arr) (OpenExp env aenv dim -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv dim
ix)
    LinearIndex ArrayVar aenv (Array dim t)
arr OpenExp env aenv Int
ix    -> Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
ppF2 (Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"!!"         Fixity
Infix Associativity
L Int
9) (ArrayVar aenv (Array dim t) -> Context -> Adoc
forall a. ArrayVar aenv a -> Context -> Adoc
ppA ArrayVar aenv (Array dim t)
arr) (OpenExp env aenv Int -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv Int
ix)
    Coerce ScalarType a
_ ScalarType t
tp OpenExp env aenv a
x         -> Operator -> (Context -> Adoc) -> Adoc
ppF1 (Adoc -> Fixity -> Associativity -> Int -> Operator
Operator (ScalarType t -> Adoc -> Adoc
withTypeRep ScalarType t
tp Adoc
"coerce") Fixity
App Associativity
L Int
10) (OpenExp env aenv a -> Context -> Adoc
forall e. OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv a
x)
    Undef ScalarType t
tp              -> ScalarType t -> Adoc -> Adoc
withTypeRep ScalarType t
tp Adoc
"undef"

  where
    ppE :: OpenExp env aenv e -> Context -> Adoc
    ppE :: OpenExp env aenv e -> Context -> Adoc
ppE OpenExp env aenv e
e Context
c = Context -> Val env -> Val aenv -> OpenExp env aenv e -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
c Val env
env Val aenv
aenv OpenExp env aenv e
e

    ppA :: ArrayVar aenv a -> Context -> Adoc
    ppA :: ArrayVar aenv a -> Context -> Adoc
ppA ArrayVar aenv a
acc Context
_ = Val aenv -> ArrayVar aenv a -> Adoc
forall aenv a. Val aenv -> ArrayVar aenv a -> Adoc
prettyArrayVar Val aenv
aenv ArrayVar aenv a
acc

    ppF :: OpenFun env aenv f -> Context -> Adoc
    ppF :: OpenFun env aenv f -> Context -> Adoc
ppF OpenFun env aenv f
f Context
_ = Adoc -> Adoc
forall ann. Doc ann -> Doc ann
parens (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ Val env -> Val aenv -> OpenFun env aenv f -> Adoc
forall env aenv f.
Val env -> Val aenv -> OpenFun env aenv f -> Adoc
prettyOpenFun Val env
env Val aenv
aenv OpenFun env aenv f
f

    ppF1 :: Operator -> (Context -> Adoc) -> Adoc
    ppF1 :: Operator -> (Context -> Adoc) -> Adoc
ppF1 Operator
op Context -> Adoc
x
      = Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Operator -> Bool
needsParens Context
ctx Operator
op)
      (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ [Adoc] -> Adoc
combine [ Operator -> Adoc
opName Operator
op, Context -> Adoc
x Context
ctx' ]
      where
        ctx' :: Context
ctx'    = Operator -> Bool
isPrefix Operator
op Bool -> (Context, Context) -> Context
forall a. Bool -> (a, a) -> a
? (Operator -> Associativity -> Context
arg Operator
op Associativity
R, Context
app)
        combine :: [Adoc] -> Adoc
combine = Operator -> Bool
isPrefix Operator
op Bool -> ([Adoc] -> Adoc, [Adoc] -> Adoc) -> [Adoc] -> Adoc
forall a. Bool -> (a, a) -> a
? ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
cat, Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
2 (Adoc -> Adoc) -> ([Adoc] -> Adoc) -> [Adoc] -> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep)

    ppF2 :: Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
    ppF2 :: Operator -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
ppF2 Operator
op Context -> Adoc
x Context -> Adoc
y
      = Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Operator -> Bool
needsParens Context
ctx Operator
op)
      (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ if Operator -> Bool
isInfix Operator
op
          then [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Context -> Adoc
x (Operator -> Associativity -> Context
arg Operator
op Associativity
L), Adoc -> Adoc
forall ann. Doc ann -> Doc ann
group ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [Operator -> Adoc
opName Operator
op, Context -> Adoc
y (Operator -> Associativity -> Context
arg Operator
op Associativity
R)]) ]
          else Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
2 (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Operator -> Adoc
opName Operator
op, Context -> Adoc
x Context
app, Context -> Adoc
y Context
app ]

    ppF3 :: Operator -> (Context -> Adoc) -> (Context -> Adoc) -> (Context -> Adoc) -> Adoc
    ppF3 :: Operator
-> (Context -> Adoc)
-> (Context -> Adoc)
-> (Context -> Adoc)
-> Adoc
ppF3 Operator
op Context -> Adoc
x Context -> Adoc
y Context -> Adoc
z
      = Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Operator -> Bool
needsParens Context
ctx Operator
op)
      (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
2
      (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Operator -> Adoc
opName Operator
op, Context -> Adoc
x Context
app, Context -> Adoc
y Context
app, Context -> Adoc
z Context
app ]

    withTypeRep :: ScalarType t -> Adoc -> Adoc
    withTypeRep :: ScalarType t -> Adoc -> Adoc
withTypeRep ScalarType t
t Adoc
op = Adoc
op Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
"@" Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> String -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty (ScalarType t -> String
forall a. Show a => a -> String
show ScalarType t
t)

prettyArrayVar
    :: forall aenv a.
       Val aenv
    -> ArrayVar aenv a
    -> Adoc
prettyArrayVar :: Val aenv -> ArrayVar aenv a -> Adoc
prettyArrayVar Val aenv
aenv (Var ArrayR a
_ Idx aenv a
idx) = Idx aenv a -> Val aenv -> Adoc
forall env t. Idx env t -> Val env -> Adoc
prj Idx aenv a
idx Val aenv
aenv

prettyLet
    :: forall env aenv t.
       Context
    -> Val env
    -> Val aenv
    -> OpenExp env aenv t
    -> Adoc
prettyLet :: Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyLet Context
ctx Val env
env0 Val aenv
aenv
  = Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Operator -> Bool
needsParens Context
ctx Operator
"let")
  (Adoc -> Adoc)
-> (OpenExp env aenv t -> Adoc) -> OpenExp env aenv t -> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Adoc -> Adoc
forall ann. Doc ann -> Doc ann
align (Adoc -> Adoc)
-> (OpenExp env aenv t -> Adoc) -> OpenExp env aenv t -> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Adoc], Adoc) -> Adoc
wrap (([Adoc], Adoc) -> Adoc)
-> (OpenExp env aenv t -> ([Adoc], Adoc))
-> OpenExp env aenv t
-> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Val env -> OpenExp env aenv t -> ([Adoc], Adoc)
forall env' e. Val env' -> OpenExp env' aenv e -> ([Adoc], Adoc)
collect Val env
env0
  where
    collect :: Val env' -> OpenExp env' aenv e -> ([Adoc], Adoc)
    collect :: Val env' -> OpenExp env' aenv e -> ([Adoc], Adoc)
collect Val env'
env =
      \case
        Let ELeftHandSide bnd_t env' env'
lhs OpenExp env' aenv bnd_t
e1 OpenExp env' aenv e
e2 ->
          let (Val env'
env', Adoc
v)       = Bool
-> Val env' -> ELeftHandSide bnd_t env' env' -> (Val env', Adoc)
forall env (s :: * -> *) arrs env'.
Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyELhs Bool
False Val env'
env ELeftHandSide bnd_t env' env'
lhs
              e1' :: Adoc
e1'             = Val env' -> OpenExp env' aenv bnd_t -> Adoc
forall env' t'. Val env' -> OpenExp env' aenv t' -> Adoc
ppE Val env'
env OpenExp env' aenv bnd_t
e1
              bnd :: Adoc
bnd | OpenExp env' aenv bnd_t -> Bool
forall env' t'. OpenExp env' aenv t' -> Bool
isLet OpenExp env' aenv bnd_t
e1  = Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
nest Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vsep [Adoc
v Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
forall ann. Doc ann
equals, Adoc
e1'])
                  | Bool
otherwise = Adoc
v Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc -> Adoc
forall ann. Doc ann -> Doc ann
align (Adoc
forall ann. Doc ann
equals Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
e1')
              ([Adoc]
bnds, Adoc
body)    = Val env' -> OpenExp env' aenv e -> ([Adoc], Adoc)
forall env' e. Val env' -> OpenExp env' aenv e -> ([Adoc], Adoc)
collect Val env'
env' OpenExp env' aenv e
e2
          in
          (Adoc
bndAdoc -> [Adoc] -> [Adoc]
forall a. a -> [a] -> [a]
:[Adoc]
bnds, Adoc
body)
        --
        OpenExp env' aenv e
next     -> ([], Val env' -> OpenExp env' aenv e -> Adoc
forall env' t'. Val env' -> OpenExp env' aenv t' -> Adoc
ppE Val env'
env OpenExp env' aenv e
next)

    isLet :: OpenExp env' aenv t' -> Bool
    isLet :: OpenExp env' aenv t' -> Bool
isLet Let{} = Bool
True
    isLet OpenExp env' aenv t'
_     = Bool
False

    ppE :: Val env' -> OpenExp env' aenv t' -> Adoc
    ppE :: Val env' -> OpenExp env' aenv t' -> Adoc
ppE Val env'
env = Context -> Val env' -> Val aenv -> OpenExp env' aenv t' -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
context0 Val env'
env Val aenv
aenv

    wrap :: ([Adoc], Adoc) -> Adoc
    wrap :: ([Adoc], Adoc) -> Adoc
wrap ([],   Adoc
body) = Adoc
body  -- shouldn't happen!
    wrap ([Adoc
b],  Adoc
body)
      = [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [ Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
nest Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [Adoc
let_, Adoc
b]), Adoc
in_, Adoc
body ]
    wrap ([Adoc]
bnds, Adoc
body)
      = [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vsep [ Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
nest Int
shiftwidth ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vsep (Adoc
let_ Adoc -> [Adoc] -> [Adoc]
forall a. a -> [a] -> [a]
: [Adoc]
bnds))
             , Adoc
in_
             , Adoc
body
             ]

prettyTuple
    :: forall env aenv t.
       Context
    -> Val env
    -> Val aenv
    -> OpenExp env aenv t
    -> Adoc
prettyTuple :: Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyTuple Context
ctx Val env
env Val aenv
aenv OpenExp env aenv t
exp = case OpenExp env aenv t -> Maybe [Adoc]
forall t'. OpenExp env aenv t' -> Maybe [Adoc]
collect OpenExp env aenv t
exp of
    Maybe [Adoc]
Nothing  -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann
align (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ OpenExp env aenv t -> Adoc
forall t'. OpenExp env aenv t' -> Adoc
ppPair OpenExp env aenv t
exp
    Just [Adoc]
tup ->
      case [Adoc]
tup of
        []  -> Adoc
"()"
        [Adoc
t] -> Adoc
t
        [Adoc]
_   -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann
align (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Int
ctxPrecedence Context
ctx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Adoc
"T" Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Int -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty ([Adoc] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Adoc]
tup) Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [Adoc]
tup)
  where
    ppPair :: OpenExp env aenv t' -> Adoc
    ppPair :: OpenExp env aenv t' -> Adoc
ppPair (Pair OpenExp env aenv t1
e1 OpenExp env aenv t2
e2) = Adoc
"(" Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> OpenExp env aenv t1 -> Adoc
forall t'. OpenExp env aenv t' -> Adoc
ppPair OpenExp env aenv t1
e1 Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
"," Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Context -> Val env -> Val aenv -> OpenExp env aenv t2 -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
context0 Val env
env Val aenv
aenv OpenExp env aenv t2
e2 Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Adoc
")"
    ppPair OpenExp env aenv t'
e            = Context -> Val env -> Val aenv -> OpenExp env aenv t' -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
context0 Val env
env Val aenv
aenv OpenExp env aenv t'
e

    collect :: OpenExp env aenv t' -> Maybe [Adoc]
    collect :: OpenExp env aenv t' -> Maybe [Adoc]
collect OpenExp env aenv t'
Nil                = [Adoc] -> Maybe [Adoc]
forall a. a -> Maybe a
Just []
    collect (Pair OpenExp env aenv t1
e1 OpenExp env aenv t2
e2)
      | Just [Adoc]
tup <- OpenExp env aenv t1 -> Maybe [Adoc]
forall t'. OpenExp env aenv t' -> Maybe [Adoc]
collect OpenExp env aenv t1
e1 = [Adoc] -> Maybe [Adoc]
forall a. a -> Maybe a
Just ([Adoc] -> Maybe [Adoc]) -> [Adoc] -> Maybe [Adoc]
forall a b. (a -> b) -> a -> b
$ [Adoc]
tup [Adoc] -> [Adoc] -> [Adoc]
forall a. [a] -> [a] -> [a]
++ [Context -> Val env -> Val aenv -> OpenExp env aenv t2 -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
app Val env
env Val aenv
aenv OpenExp env aenv t2
e2]
    collect OpenExp env aenv t'
_                  = Maybe [Adoc]
forall a. Maybe a
Nothing

prettyCase
    :: Val env
    -> Val aenv
    -> OpenExp env aenv a
    -> [(TAG, OpenExp env aenv b)]
    -> Maybe (OpenExp env aenv b)
    -> Adoc
prettyCase :: Val env
-> Val aenv
-> OpenExp env aenv a
-> [(PrimBool, OpenExp env aenv b)]
-> Maybe (OpenExp env aenv b)
-> Adoc
prettyCase Val env
env Val aenv
aenv OpenExp env aenv a
x [(PrimBool, OpenExp env aenv b)]
xs Maybe (OpenExp env aenv b)
def
  = Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
hang Int
shiftwidth
  (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vsep [ Adoc
case_ Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
x' Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
of_
         , Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
flatAlt ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
vcat [Adoc]
xs') (Adoc -> Adoc -> Adoc -> [Adoc] -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann -> [Doc ann] -> Doc ann
encloseSep Adoc
"{ " Adoc
" }" Adoc
"; " [Adoc]
xs')
         ]
  where
    x' :: Adoc
x'  = Context -> Val env -> Val aenv -> OpenExp env aenv a -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
context0 Val env
env Val aenv
aenv OpenExp env aenv a
x
    xs' :: [Adoc]
xs' = ((PrimBool, OpenExp env aenv b) -> Adoc)
-> [(PrimBool, OpenExp env aenv b)] -> [Adoc]
forall a b. (a -> b) -> [a] -> [b]
map (\(PrimBool
t,OpenExp env aenv b
e) -> PrimBool -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty PrimBool
t Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
"->" Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Context -> Val env -> Val aenv -> OpenExp env aenv b -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
context0 Val env
env Val aenv
aenv OpenExp env aenv b
e) [(PrimBool, OpenExp env aenv b)]
xs
       [Adoc] -> [Adoc] -> [Adoc]
forall a. [a] -> [a] -> [a]
++ case Maybe (OpenExp env aenv b)
def of
            Maybe (OpenExp env aenv b)
Nothing -> []
            Just OpenExp env aenv b
d  -> [Adoc
"_" Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
"->" Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Context -> Val env -> Val aenv -> OpenExp env aenv b -> Adoc
forall env aenv t.
Context -> Val env -> Val aenv -> OpenExp env aenv t -> Adoc
prettyOpenExp Context
context0 Val env
env Val aenv
aenv OpenExp env aenv b
d]

{-

prettyAtuple
    :: forall acc aenv arrs.
       PrettyAcc acc
    -> ExtractAcc acc
    -> Val aenv
    -> PreOpenAcc acc aenv arrs
    -> Adoc
prettyAtuple prettyAcc extractAcc aenv0 acc = case collect acc of
  Just tup -> align $ "T" <> pretty (length tup) <+> sep tup
  Nothing  -> align $ ppPair acc
  where
    ppPair :: PreOpenAcc acc aenv arrs' -> Adoc
    ppPair (Apair a1 a2) = "(" <> ppPair (extractAcc a1) <> "," <+> prettyAcc context0 aenv0 a2 <> ")"
    ppPair a             = prettyPreOpenAcc context0 prettyAcc extractAcc aenv0 a

    collect :: PreOpenAcc acc aenv arrs' -> Maybe [Adoc]
    collect Anil          = Just []
    collect (Apair a1 a2)
      | Just tup <- collect $ extractAcc a1
                          = Just $ tup ++ [prettyAcc app aenv0 a2]
    collect _             = Nothing
-}

prettyConst :: TypeR e -> e -> Adoc
prettyConst :: TypeR e -> e -> Adoc
prettyConst TypeR e
tp e
x =
  let y :: String
y = TypeR e -> e -> String
forall e. TypeR e -> e -> String
showElt TypeR e
tp e
x
  in  Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf ((Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Char -> Bool
isSpace String
y) (String -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty String
y)

prettyPrimConst :: PrimConst a -> Adoc
prettyPrimConst :: PrimConst a -> Adoc
prettyPrimConst PrimMinBound{} = Adoc
"minBound"
prettyPrimConst PrimMaxBound{} = Adoc
"maxBound"
prettyPrimConst PrimPi{}       = Adoc
"pi"


-- Primitive operators
-- -------------------
--
-- The core of the pretty printer is how to correctly handle precedence,
-- associativity, and fixity of the primitive scalar operators.
--

data Direction = L | N | R
  deriving Associativity -> Associativity -> Bool
(Associativity -> Associativity -> Bool)
-> (Associativity -> Associativity -> Bool) -> Eq Associativity
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Associativity -> Associativity -> Bool
$c/= :: Associativity -> Associativity -> Bool
== :: Associativity -> Associativity -> Bool
$c== :: Associativity -> Associativity -> Bool
Eq

data Fixity = App | Infix | Prefix
  deriving Fixity -> Fixity -> Bool
(Fixity -> Fixity -> Bool)
-> (Fixity -> Fixity -> Bool) -> Eq Fixity
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Fixity -> Fixity -> Bool
$c/= :: Fixity -> Fixity -> Bool
== :: Fixity -> Fixity -> Bool
$c== :: Fixity -> Fixity -> Bool
Eq

type Precedence    = Int
type Associativity = Direction

data Context = Context
  { Context -> Associativity
ctxAssociativity  :: Associativity
  , Context -> Associativity
ctxPosition       :: Direction
  , Context -> Int
ctxPrecedence     :: Precedence
  }

data Operator = Operator
  { Operator -> Adoc
opName            :: Adoc
  , Operator -> Fixity
opFixity          :: Fixity
  , Operator -> Associativity
opAssociativity   :: Associativity
  , Operator -> Int
opPrecedence      :: Precedence
  }

instance IsString Operator where
  fromString :: String -> Operator
fromString String
s = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator (String -> Adoc
forall a. IsString a => String -> a
fromString String
s) Fixity
App Associativity
L Int
10

needsParens :: Context -> Operator -> Bool
needsParens :: Context -> Operator -> Bool
needsParens Context{Int
Associativity
ctxPrecedence :: Int
ctxPosition :: Associativity
ctxAssociativity :: Associativity
ctxPosition :: Context -> Associativity
ctxAssociativity :: Context -> Associativity
ctxPrecedence :: Context -> Int
..} Operator{Int
Adoc
Fixity
Associativity
opPrecedence :: Int
opAssociativity :: Associativity
opFixity :: Fixity
opName :: Adoc
opPrecedence :: Operator -> Int
opAssociativity :: Operator -> Associativity
opFixity :: Operator -> Fixity
opName :: Operator -> Adoc
..}
  | Int
ctxPrecedence     Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
opPrecedence    = Bool
False
  | Int
ctxPrecedence     Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
opPrecedence    = Bool
True
  | Associativity
ctxAssociativity Associativity -> Associativity -> Bool
forall a. Eq a => a -> a -> Bool
/= Associativity
opAssociativity = Bool
True
  | Bool
otherwise                           = Associativity
ctxPosition Associativity -> Associativity -> Bool
forall a. Eq a => a -> a -> Bool
/= Associativity
opAssociativity

context0 :: Context
context0 :: Context
context0 = Associativity -> Associativity -> Int -> Context
Context Associativity
N Associativity
N Int
0

app :: Context
app :: Context
app = Associativity -> Associativity -> Int -> Context
Context Associativity
L Associativity
N Int
10

arg :: Operator -> Direction -> Context
arg :: Operator -> Associativity -> Context
arg Operator{Int
Adoc
Fixity
Associativity
opPrecedence :: Int
opAssociativity :: Associativity
opFixity :: Fixity
opName :: Adoc
opPrecedence :: Operator -> Int
opAssociativity :: Operator -> Associativity
opFixity :: Operator -> Fixity
opName :: Operator -> Adoc
..} Associativity
side = Associativity -> Associativity -> Int -> Context
Context Associativity
opAssociativity Associativity
side Int
opPrecedence

isPrefix :: Operator -> Bool
isPrefix :: Operator -> Bool
isPrefix Operator{Int
Adoc
Fixity
Associativity
opPrecedence :: Int
opAssociativity :: Associativity
opFixity :: Fixity
opName :: Adoc
opPrecedence :: Operator -> Int
opAssociativity :: Operator -> Associativity
opFixity :: Operator -> Fixity
opName :: Operator -> Adoc
..} = Fixity
opFixity Fixity -> Fixity -> Bool
forall a. Eq a => a -> a -> Bool
== Fixity
Prefix

isInfix :: Operator -> Bool
isInfix :: Operator -> Bool
isInfix Operator{Int
Adoc
Fixity
Associativity
opPrecedence :: Int
opAssociativity :: Associativity
opFixity :: Fixity
opName :: Adoc
opPrecedence :: Operator -> Int
opAssociativity :: Operator -> Associativity
opFixity :: Operator -> Fixity
opName :: Operator -> Adoc
..}  = Fixity
opFixity Fixity -> Fixity -> Bool
forall a. Eq a => a -> a -> Bool
== Fixity
Infix

primOperator :: PrimFun a -> Operator
primOperator :: PrimFun a -> Operator
primOperator PrimAdd{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator (Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'+')         Fixity
Infix  Associativity
L Int
6
primOperator PrimSub{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator (Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'-')         Fixity
Infix  Associativity
L Int
6
primOperator PrimMul{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator (Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'*')         Fixity
Infix  Associativity
L Int
7
primOperator PrimNeg{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator (Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'-')         Fixity
Prefix Associativity
L Int
6  -- Haskell's only prefix operator
primOperator PrimAbs{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"abs"                Fixity
App    Associativity
L Int
10
primOperator PrimSig{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"signum"             Fixity
App    Associativity
L Int
10
primOperator PrimQuot{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"quot"               Fixity
App    Associativity
L Int
10
primOperator PrimRem{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"rem"                Fixity
App    Associativity
L Int
10
primOperator PrimQuotRem{}            = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"quotRem"            Fixity
App    Associativity
L Int
10
primOperator PrimIDiv{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"div"                Fixity
App    Associativity
L Int
10
primOperator PrimMod{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"mod"                Fixity
App    Associativity
L Int
10
primOperator PrimDivMod{}             = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"divMod"             Fixity
App    Associativity
L Int
10
primOperator PrimBAnd{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
".&."                Fixity
Infix  Associativity
L Int
7
primOperator PrimBOr{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
".|."                Fixity
Infix  Associativity
L Int
5
primOperator PrimBXor{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"xor"                Fixity
App    Associativity
L Int
10
primOperator PrimBNot{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"complement"         Fixity
App    Associativity
L Int
10
primOperator PrimBShiftL{}            = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"shiftL"             Fixity
App    Associativity
L Int
10
primOperator PrimBShiftR{}            = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"shiftR"             Fixity
App    Associativity
L Int
10
primOperator PrimBRotateL{}           = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"rotateL"            Fixity
App    Associativity
L Int
10
primOperator PrimBRotateR{}           = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"rotateR"            Fixity
App    Associativity
L Int
10
primOperator PrimPopCount{}           = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"popCount"           Fixity
App    Associativity
L Int
10
primOperator PrimCountLeadingZeros{}  = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"countLeadingZeros"  Fixity
App    Associativity
L Int
10
primOperator PrimCountTrailingZeros{} = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"countTrailingZeros" Fixity
App    Associativity
L Int
10
primOperator PrimFDiv{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator (Char -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Char
'/')         Fixity
Infix  Associativity
L Int
7
primOperator PrimRecip{}              = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"recip"              Fixity
App    Associativity
L Int
10
primOperator PrimSin{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"sin"                Fixity
App    Associativity
L Int
10
primOperator PrimCos{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"cos"                Fixity
App    Associativity
L Int
10
primOperator PrimTan{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"tan"                Fixity
App    Associativity
L Int
10
primOperator PrimAsin{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"asin"               Fixity
App    Associativity
L Int
10
primOperator PrimAcos{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"acos"               Fixity
App    Associativity
L Int
10
primOperator PrimAtan{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"atan"               Fixity
App    Associativity
L Int
10
primOperator PrimSinh{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"sinh"               Fixity
App    Associativity
L Int
10
primOperator PrimCosh{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"cosh"               Fixity
App    Associativity
L Int
10
primOperator PrimTanh{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"tanh"               Fixity
App    Associativity
L Int
10
primOperator PrimAsinh{}              = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"asinh"              Fixity
App    Associativity
L Int
10
primOperator PrimAcosh{}              = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"acosh"              Fixity
App    Associativity
L Int
10
primOperator PrimAtanh{}              = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"atanh"              Fixity
App    Associativity
L Int
10
primOperator PrimExpFloating{}        = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"exp"                Fixity
App    Associativity
L Int
10
primOperator PrimSqrt{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"sqrt"               Fixity
App    Associativity
L Int
10
primOperator PrimLog{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"log"                Fixity
App    Associativity
L Int
10
primOperator PrimFPow{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"**"                 Fixity
Infix  Associativity
R Int
8
primOperator PrimLogBase{}            = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"logBase"            Fixity
App    Associativity
L Int
10
primOperator PrimTruncate{}           = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"truncate"           Fixity
App    Associativity
L Int
10
primOperator PrimRound{}              = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"round"              Fixity
App    Associativity
L Int
10
primOperator PrimFloor{}              = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"floor"              Fixity
App    Associativity
L Int
10
primOperator PrimCeiling{}            = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"ceiling"            Fixity
App    Associativity
L Int
10
primOperator PrimAtan2{}              = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"atan2"              Fixity
App    Associativity
L Int
10
primOperator PrimIsNaN{}              = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"isNaN"              Fixity
App    Associativity
L Int
10
primOperator PrimIsInfinite{}         = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"isInfinite"         Fixity
App    Associativity
L Int
10
primOperator PrimLt{}                 = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"<"                  Fixity
Infix  Associativity
N Int
4
primOperator PrimGt{}                 = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
">"                  Fixity
Infix  Associativity
N Int
4
primOperator PrimLtEq{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"<="                 Fixity
Infix  Associativity
N Int
4
primOperator PrimGtEq{}               = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
">="                 Fixity
Infix  Associativity
N Int
4
primOperator PrimEq{}                 = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"=="                 Fixity
Infix  Associativity
N Int
4
primOperator PrimNEq{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"/="                 Fixity
Infix  Associativity
N Int
4
primOperator PrimMax{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"max"                Fixity
App    Associativity
L Int
10
primOperator PrimMin{}                = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"min"                Fixity
App    Associativity
L Int
10
primOperator PrimFun a
PrimLAnd                 = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"&&"                 Fixity
Infix  Associativity
R Int
3
primOperator PrimFun a
PrimLOr                  = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"||"                 Fixity
Infix  Associativity
R Int
2
primOperator PrimFun a
PrimLNot                 = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"not"                Fixity
App    Associativity
L Int
10
primOperator PrimFromIntegral{}       = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"fromIntegral"       Fixity
App    Associativity
L Int
10
primOperator PrimToFloating{}         = Adoc -> Fixity -> Associativity -> Int -> Operator
Operator Adoc
"toFloating"         Fixity
App    Associativity
L Int
10


-- Environments
-- ------------

data Val env where
  Empty ::                    Val ()
  Push  :: Val env -> Adoc -> Val (env, t)

class PrettyEnv env where
  prettyEnv :: Adoc -> Val env

instance PrettyEnv () where
  prettyEnv :: Adoc -> Val ()
prettyEnv Adoc
_ = Val ()
Empty

instance PrettyEnv env => PrettyEnv (env, t) where
  prettyEnv :: Adoc -> Val (env, t)
prettyEnv Adoc
v =
    let env :: Val env
env = Adoc -> Val env
forall env. PrettyEnv env => Adoc -> Val env
prettyEnv Adoc
v :: Val env
        x :: Adoc
x   = Adoc
v Adoc -> Adoc -> Adoc
forall a. Semigroup a => a -> a -> a
<> Int -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty (Val env -> Int
forall env. Val env -> Int
sizeEnv Val env
env)
    in
    Val env
env Val env -> Adoc -> Val (env, t)
forall env t. Val env -> Adoc -> Val (env, t)
`Push` Adoc
x

sizeEnv :: Val env -> Int
sizeEnv :: Val env -> Int
sizeEnv Val env
Empty        = Int
0
sizeEnv (Push Val env
env Adoc
_) = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Val env -> Int
forall env. Val env -> Int
sizeEnv Val env
env

prj :: Idx env t -> Val env -> Adoc
prj :: Idx env t -> Val env -> Adoc
prj Idx env t
ZeroIdx      (Push Val env
_ Adoc
v)   = Adoc
v
prj (SuccIdx Idx env t
ix) (Push Val env
env Adoc
_) = Idx env t -> Val env -> Adoc
forall env t. Idx env t -> Val env -> Adoc
prj Idx env t
ix Val env
Val env
env


-- Utilities
-- ---------

shiftwidth :: Int
shiftwidth :: Int
shiftwidth = Int
2

infix 0 ?
(?) :: Bool -> (a, a) -> a
Bool
True  ? :: Bool -> (a, a) -> a
? (a
t,a
_) = a
t
Bool
False ? (a
_,a
f) = a
f

parensIf :: Bool -> Doc ann -> Doc ann
parensIf :: Bool -> Doc ann -> Doc ann
parensIf Bool
True  = Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
group (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
align
parensIf Bool
False = Doc ann -> Doc ann
forall a. a -> a
id