{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

-- | Futhark prettyprinter.  This module defines 'Pretty' instances
-- for the AST defined in "Futhark.IR.Syntax",
-- but also a number of convenience functions if you don't want to use
-- the interface from 'Pretty'.
module Futhark.IR.Pretty
  ( prettyTuple,
    prettyTupleLines,
    pretty,
    PrettyRep (..),
    ppTuple',
  )
where

import Data.Foldable (toList)
import Data.Maybe
import Futhark.IR.Syntax
import Futhark.Util.Pretty

-- | The class of representations whose annotations can be prettyprinted.
class
  ( RepTypes rep,
    Pretty (RetType rep),
    Pretty (BranchType rep),
    Pretty (FParamInfo rep),
    Pretty (LParamInfo rep),
    Pretty (LetDec rep),
    Pretty (Op rep)
  ) =>
  PrettyRep rep
  where
  ppExpDec :: ExpDec rep -> Exp rep -> Maybe Doc
  ppExpDec ExpDec rep
_ Exp rep
_ = Maybe Doc
forall a. Maybe a
Nothing

instance Pretty VName where
  ppr :: VName -> Doc
ppr (VName Name
vn Int
i) = Name -> Doc
forall a. Pretty a => a -> Doc
ppr Name
vn Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
"_" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text (Int -> String
forall a. Show a => a -> String
show Int
i)

instance Pretty Commutativity where
  ppr :: Commutativity -> Doc
ppr Commutativity
Commutative = String -> Doc
text String
"commutative"
  ppr Commutativity
Noncommutative = String -> Doc
text String
"noncommutative"

instance Pretty NoUniqueness where
  ppr :: NoUniqueness -> Doc
ppr NoUniqueness
_ = Doc
forall a. Monoid a => a
mempty

instance Pretty Shape where
  ppr :: Shape -> Doc
ppr = [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ([Doc] -> Doc) -> (Shape -> [Doc]) -> Shape -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp -> Doc) -> [SubExp] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
brackets (Doc -> Doc) -> (SubExp -> Doc) -> SubExp -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr) ([SubExp] -> [Doc]) -> (Shape -> [SubExp]) -> Shape -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shape -> [SubExp]
forall d. ShapeBase d -> [d]
shapeDims

instance Pretty a => Pretty (Ext a) where
  ppr :: Ext a -> Doc
ppr (Free a
e) = a -> Doc
forall a. Pretty a => a -> Doc
ppr a
e
  ppr (Ext Int
x) = String -> Doc
text String
"?" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text (Int -> String
forall a. Show a => a -> String
show Int
x)

instance Pretty ExtShape where
  ppr :: ExtShape -> Doc
ppr = [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ([Doc] -> Doc) -> (ExtShape -> [Doc]) -> ExtShape -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ExtSize -> Doc) -> [ExtSize] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
brackets (Doc -> Doc) -> (ExtSize -> Doc) -> ExtSize -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtSize -> Doc
forall a. Pretty a => a -> Doc
ppr) ([ExtSize] -> [Doc])
-> (ExtShape -> [ExtSize]) -> ExtShape -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtShape -> [ExtSize]
forall d. ShapeBase d -> [d]
shapeDims

instance Pretty Space where
  ppr :: Space -> Doc
ppr Space
DefaultSpace = Doc
forall a. Monoid a => a
mempty
  ppr (Space String
s) = String -> Doc
text String
"@" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
s
  ppr (ScalarSpace [SubExp]
d PrimType
t) = String -> Doc
text String
"@" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ((SubExp -> Doc) -> [SubExp] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
brackets (Doc -> Doc) -> (SubExp -> Doc) -> SubExp -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr) [SubExp]
d) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t

instance Pretty u => Pretty (TypeBase Shape u) where
  ppr :: TypeBase Shape u -> Doc
ppr (Prim PrimType
t) = PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t
  ppr (Acc VName
acc Shape
ispace [Type]
ts u
u) =
    u -> Doc
forall a. Pretty a => a -> Doc
ppr u
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
"acc" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
acc, Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
ispace, [Type] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [Type]
ts]
  ppr (Array PrimType
et (Shape [SubExp]
ds) u
u) =
    u -> Doc
forall a. Pretty a => a -> Doc
ppr u
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ((SubExp -> Doc) -> [SubExp] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
brackets (Doc -> Doc) -> (SubExp -> Doc) -> SubExp -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr) [SubExp]
ds) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
et
  ppr (Mem Space
s) = String -> Doc
text String
"mem" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
s

instance Pretty u => Pretty (TypeBase ExtShape u) where
  ppr :: TypeBase ExtShape u -> Doc
ppr (Prim PrimType
t) = PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t
  ppr (Acc VName
acc Shape
ispace [Type]
ts u
u) =
    u -> Doc
forall a. Pretty a => a -> Doc
ppr u
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
"acc" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
acc, Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
ispace, [Type] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [Type]
ts]
  ppr (Array PrimType
et (Shape [ExtSize]
ds) u
u) =
    u -> Doc
forall a. Pretty a => a -> Doc
ppr u
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ((ExtSize -> Doc) -> [ExtSize] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
brackets (Doc -> Doc) -> (ExtSize -> Doc) -> ExtSize -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExtSize -> Doc
forall a. Pretty a => a -> Doc
ppr) [ExtSize]
ds) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
et
  ppr (Mem Space
s) = String -> Doc
text String
"mem" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
s

instance Pretty u => Pretty (TypeBase Rank u) where
  ppr :: TypeBase Rank u -> Doc
ppr (Prim PrimType
t) = PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t
  ppr (Acc VName
acc Shape
ispace [Type]
ts u
u) =
    u -> Doc
forall a. Pretty a => a -> Doc
ppr u
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
"acc" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
acc, Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
ispace, [Type] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [Type]
ts]
  ppr (Array PrimType
et (Rank Int
n) u
u) =
    u -> Doc
forall a. Pretty a => a -> Doc
ppr u
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat (Int -> Doc -> [Doc]
forall a. Int -> a -> [a]
replicate Int
n (Doc -> [Doc]) -> Doc -> [Doc]
forall a b. (a -> b) -> a -> b
$ Doc -> Doc
brackets Doc
forall a. Monoid a => a
mempty) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
et
  ppr (Mem Space
s) = String -> Doc
text String
"mem" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Space -> Doc
forall a. Pretty a => a -> Doc
ppr Space
s

instance Pretty Ident where
  ppr :: Ident -> Doc
ppr Ident
ident = Type -> Doc
forall a. Pretty a => a -> Doc
ppr (Ident -> Type
identType Ident
ident) Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr (Ident -> VName
identName Ident
ident)

instance Pretty SubExp where
  ppr :: SubExp -> Doc
ppr (Var VName
v) = VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v
  ppr (Constant PrimValue
v) = PrimValue -> Doc
forall a. Pretty a => a -> Doc
ppr PrimValue
v

instance Pretty Certs where
  ppr :: Certs -> Doc
ppr (Certs []) = Doc
empty
  ppr (Certs [VName]
cs) = String -> Doc
text String
"#" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
braces ([Doc] -> Doc
commasep ((VName -> Doc) -> [VName] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Doc
forall a. Pretty a => a -> Doc
ppr [VName]
cs))

instance PrettyRep rep => Pretty (Stms rep) where
  ppr :: Stms rep -> Doc
ppr = [Doc] -> Doc
stack ([Doc] -> Doc) -> (Stms rep -> [Doc]) -> Stms rep -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stm rep -> Doc) -> [Stm rep] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Stm rep -> Doc
forall a. Pretty a => a -> Doc
ppr ([Stm rep] -> [Doc])
-> (Stms rep -> [Stm rep]) -> Stms rep -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList

instance Pretty SubExpRes where
  ppr :: SubExpRes -> Doc
ppr (SubExpRes Certs
cs SubExp
se) = [Doc] -> Doc
spread ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ Certs -> [Doc]
certAnnots Certs
cs [Doc] -> [Doc] -> [Doc]
forall a. [a] -> [a] -> [a]
++ [SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
se]

instance PrettyRep rep => Pretty (Body rep) where
  ppr :: Body rep -> Doc
ppr (Body BodyDec rep
_ Stms rep
stms [SubExpRes]
res)
    | Stms rep -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Stms rep
stms = Doc -> Doc
braces ([Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (SubExpRes -> Doc) -> [SubExpRes] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> Doc
forall a. Pretty a => a -> Doc
ppr [SubExpRes]
res)
    | Bool
otherwise =
      [Doc] -> Doc
stack ((Stm rep -> Doc) -> [Stm rep] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Stm rep -> Doc
forall a. Pretty a => a -> Doc
ppr ([Stm rep] -> [Doc]) -> [Stm rep] -> [Doc]
forall a b. (a -> b) -> a -> b
$ Stms rep -> [Stm rep]
forall rep. Stms rep -> [Stm rep]
stmsToList Stms rep
stms)
        Doc -> Doc -> Doc
</> String -> Doc
text String
"in" Doc -> Doc -> Doc
<+> Doc -> Doc
braces ([Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (SubExpRes -> Doc) -> [SubExpRes] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map SubExpRes -> Doc
forall a. Pretty a => a -> Doc
ppr [SubExpRes]
res)

instance Pretty Attr where
  ppr :: Attr -> Doc
ppr (AttrName Name
v) = Name -> Doc
forall a. Pretty a => a -> Doc
ppr Name
v
  ppr (AttrInt Integer
x) = Integer -> Doc
forall a. Pretty a => a -> Doc
ppr Integer
x
  ppr (AttrComp Name
f [Attr]
attrs) = Name -> Doc
forall a. Pretty a => a -> Doc
ppr Name
f Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens ([Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (Attr -> Doc) -> [Attr] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Attr -> Doc
forall a. Pretty a => a -> Doc
ppr [Attr]
attrs)

attrAnnots :: Attrs -> [Doc]
attrAnnots :: Attrs -> [Doc]
attrAnnots = (Attr -> Doc) -> [Attr] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Attr -> Doc
forall a. Pretty a => a -> Doc
f ([Attr] -> [Doc]) -> (Attrs -> [Attr]) -> Attrs -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set Attr -> [Attr]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Set Attr -> [Attr]) -> (Attrs -> Set Attr) -> Attrs -> [Attr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Attrs -> Set Attr
unAttrs
  where
    f :: a -> Doc
f a
v = String -> Doc
text String
"#[" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> a -> Doc
forall a. Pretty a => a -> Doc
ppr a
v Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
"]"

stmAttrAnnots :: Stm rep -> [Doc]
stmAttrAnnots :: Stm rep -> [Doc]
stmAttrAnnots = Attrs -> [Doc]
attrAnnots (Attrs -> [Doc]) -> (Stm rep -> Attrs) -> Stm rep -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StmAux (ExpDec rep) -> Attrs
forall dec. StmAux dec -> Attrs
stmAuxAttrs (StmAux (ExpDec rep) -> Attrs)
-> (Stm rep -> StmAux (ExpDec rep)) -> Stm rep -> Attrs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> StmAux (ExpDec rep)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux

certAnnots :: Certs -> [Doc]
certAnnots :: Certs -> [Doc]
certAnnots Certs
cs
  | Certs
cs Certs -> Certs -> Bool
forall a. Eq a => a -> a -> Bool
== Certs
forall a. Monoid a => a
mempty = []
  | Bool
otherwise = [Certs -> Doc
forall a. Pretty a => a -> Doc
ppr Certs
cs]

stmCertAnnots :: Stm rep -> [Doc]
stmCertAnnots :: Stm rep -> [Doc]
stmCertAnnots = Certs -> [Doc]
certAnnots (Certs -> [Doc]) -> (Stm rep -> Certs) -> Stm rep -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StmAux (ExpDec rep) -> Certs
forall dec. StmAux dec -> Certs
stmAuxCerts (StmAux (ExpDec rep) -> Certs)
-> (Stm rep -> StmAux (ExpDec rep)) -> Stm rep -> Certs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> StmAux (ExpDec rep)
forall rep. Stm rep -> StmAux (ExpDec rep)
stmAux

instance Pretty Attrs where
  ppr :: Attrs -> Doc
ppr = [Doc] -> Doc
spread ([Doc] -> Doc) -> (Attrs -> [Doc]) -> Attrs -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Attrs -> [Doc]
attrAnnots

instance Pretty (PatElemT dec) => Pretty (PatT dec) where
  ppr :: PatT dec -> Doc
ppr (Pat [PatElemT dec]
xs) = Doc -> Doc
braces (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
commastack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (PatElemT dec -> Doc) -> [PatElemT dec] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT dec -> Doc
forall a. Pretty a => a -> Doc
ppr [PatElemT dec]
xs

instance Pretty t => Pretty (PatElemT t) where
  ppr :: PatElemT t -> Doc
ppr (PatElem VName
name t
t) = VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name Doc -> Doc -> Doc
<+> Doc
colon Doc -> Doc -> Doc
<+> Doc -> Doc
align (t -> Doc
forall a. Pretty a => a -> Doc
ppr t
t)

instance Pretty t => Pretty (Param t) where
  ppr :: Param t -> Doc
ppr (Param Attrs
attrs VName
name t
t) =
    [Doc] -> Doc -> Doc
annot (Attrs -> [Doc]
attrAnnots Attrs
attrs) (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
name Doc -> Doc -> Doc
<+> Doc
colon Doc -> Doc -> Doc
<+> Doc -> Doc
align (t -> Doc
forall a. Pretty a => a -> Doc
ppr t
t)

instance PrettyRep rep => Pretty (Stm rep) where
  ppr :: Stm rep -> Doc
ppr stm :: Stm rep
stm@(Let Pat rep
pat StmAux (ExpDec rep)
aux Exp rep
e) =
    Doc -> Doc
align (Doc -> Doc) -> (Doc -> Doc) -> Doc -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Doc -> Doc
hang Int
2 (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$
      String -> Doc
text String
"let" Doc -> Doc -> Doc
<+> Doc -> Doc
align (Pat rep -> Doc
forall a. Pretty a => a -> Doc
ppr Pat rep
pat)
        Doc -> Doc -> Doc
<+> case (Bool
linebreak, [Doc]
stmannot) of
          (Bool
True, []) -> Doc
equals Doc -> Doc -> Doc
</> Exp rep -> Doc
forall a. Pretty a => a -> Doc
ppr Exp rep
e
          (Bool
False, []) -> Doc
equals Doc -> Doc -> Doc
<+/> Exp rep -> Doc
forall a. Pretty a => a -> Doc
ppr Exp rep
e
          (Bool
_, [Doc]
ann) -> Doc
equals Doc -> Doc -> Doc
</> ([Doc] -> Doc
stack [Doc]
ann Doc -> Doc -> Doc
</> Exp rep -> Doc
forall a. Pretty a => a -> Doc
ppr Exp rep
e)
    where
      linebreak :: Bool
linebreak = case Exp rep
e of
        BasicOp BinOp {} -> Bool
False
        BasicOp CmpOp {} -> Bool
False
        BasicOp ConvOp {} -> Bool
False
        BasicOp UnOp {} -> Bool
False
        BasicOp SubExp {} -> Bool
False
        Exp rep
_ -> Bool
True

      stmannot :: [Doc]
stmannot =
        [[Doc]] -> [Doc]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
          [ Maybe Doc -> [Doc]
forall a. Maybe a -> [a]
maybeToList (ExpDec rep -> Exp rep -> Maybe Doc
forall rep. PrettyRep rep => ExpDec rep -> Exp rep -> Maybe Doc
ppExpDec (StmAux (ExpDec rep) -> ExpDec rep
forall dec. StmAux dec -> dec
stmAuxDec StmAux (ExpDec rep)
aux) Exp rep
e),
            Stm rep -> [Doc]
forall rep. Stm rep -> [Doc]
stmAttrAnnots Stm rep
stm,
            Stm rep -> [Doc]
forall rep. Stm rep -> [Doc]
stmCertAnnots Stm rep
stm
          ]

instance Pretty a => Pretty (Slice a) where
  ppr :: Slice a -> Doc
ppr (Slice [DimIndex a]
xs) = Doc -> Doc
brackets ([Doc] -> Doc
commasep ((DimIndex a -> Doc) -> [DimIndex a] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map DimIndex a -> Doc
forall a. Pretty a => a -> Doc
ppr [DimIndex a]
xs))

instance Pretty d => Pretty (FlatDimIndex d) where
  ppr :: FlatDimIndex d -> Doc
ppr (FlatDimIndex d
n d
s) = d -> Doc
forall a. Pretty a => a -> Doc
ppr d
n Doc -> Doc -> Doc
<+> String -> Doc
text String
":" Doc -> Doc -> Doc
<+> d -> Doc
forall a. Pretty a => a -> Doc
ppr d
s

instance Pretty a => Pretty (FlatSlice a) where
  ppr :: FlatSlice a -> Doc
ppr (FlatSlice a
offset [FlatDimIndex a]
xs) = Doc -> Doc
brackets (a -> Doc
forall a. Pretty a => a -> Doc
ppr a
offset Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
";" Doc -> Doc -> Doc
<+> [Doc] -> Doc
commasep ((FlatDimIndex a -> Doc) -> [FlatDimIndex a] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map FlatDimIndex a -> Doc
forall a. Pretty a => a -> Doc
ppr [FlatDimIndex a]
xs))

instance Pretty BasicOp where
  ppr :: BasicOp -> Doc
ppr (SubExp SubExp
se) = SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
se
  ppr (Opaque OpaqueOp
OpaqueNil SubExp
e) = String -> Doc
text String
"opaque" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
e]
  ppr (Opaque (OpaqueTrace String
s) SubExp
e) = String -> Doc
text String
"trace" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [String -> Doc
forall a. Pretty a => a -> Doc
ppr (String -> String
forall a. Show a => a -> String
show String
s), SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
e]
  ppr (ArrayLit [SubExp]
es Type
rt) =
    case Type
rt of
      Array {} -> Doc -> Doc
brackets (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
commastack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (SubExp -> Doc) -> [SubExp] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr [SubExp]
es
      Type
_ -> Doc -> Doc
brackets (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (SubExp -> Doc) -> [SubExp] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr [SubExp]
es
      Doc -> Doc -> Doc
<+> Doc
colon
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"[]" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Type -> Doc
forall a. Pretty a => a -> Doc
ppr Type
rt
  ppr (BinOp BinOp
bop SubExp
x SubExp
y) = BinOp -> Doc
forall a. Pretty a => a -> Doc
ppr BinOp
bop Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens (SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
x Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma Doc -> Doc -> Doc
<+> SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
y)
  ppr (CmpOp CmpOp
op SubExp
x SubExp
y) = CmpOp -> Doc
forall a. Pretty a => a -> Doc
ppr CmpOp
op Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens (SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
x Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma Doc -> Doc -> Doc
<+> SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
y)
  ppr (ConvOp ConvOp
conv SubExp
x) =
    String -> Doc
text (ConvOp -> String
convOpFun ConvOp
conv) Doc -> Doc -> Doc
<+> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
fromtype Doc -> Doc -> Doc
<+> SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
x Doc -> Doc -> Doc
<+> String -> Doc
text String
"to" Doc -> Doc -> Doc
<+> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
totype
    where
      (PrimType
fromtype, PrimType
totype) = ConvOp -> (PrimType, PrimType)
convOpType ConvOp
conv
  ppr (UnOp UnOp
op SubExp
e) = UnOp -> Doc
forall a. Pretty a => a -> Doc
ppr UnOp
op Doc -> Doc -> Doc
<+> Int -> SubExp -> Doc
forall a. Pretty a => Int -> a -> Doc
pprPrec Int
9 SubExp
e
  ppr (Index VName
v Slice SubExp
slice) = VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Slice SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr Slice SubExp
slice
  ppr (Update Safety
safety VName
src Slice SubExp
slice SubExp
se) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
src Doc -> Doc -> Doc
<+> Doc
with Doc -> Doc -> Doc
<+> Slice SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr Slice SubExp
slice Doc -> Doc -> Doc
<+> String -> Doc
text String
"=" Doc -> Doc -> Doc
<+> SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
se
    where
      with :: Doc
with = case Safety
safety of
        Safety
Unsafe -> String -> Doc
text String
"with"
        Safety
Safe -> String -> Doc
text String
"with?"
  ppr (FlatIndex VName
v FlatSlice SubExp
slice) = VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
v Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> FlatSlice SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr FlatSlice SubExp
slice
  ppr (FlatUpdate VName
src FlatSlice SubExp
slice VName
se) =
    VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
src Doc -> Doc -> Doc
<+> String -> Doc
text String
"with" Doc -> Doc -> Doc
<+> FlatSlice SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr FlatSlice SubExp
slice Doc -> Doc -> Doc
<+> String -> Doc
text String
"=" Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
se
  ppr (Iota SubExp
e SubExp
x SubExp
s IntType
et) = String -> Doc
text String
"iota" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
et' Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
e, SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
x, SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
s]
    where
      et' :: Doc
et' = String -> Doc
text (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
forall a b. (a -> b) -> a -> b
$ PrimType -> Int
primBitSize (PrimType -> Int) -> PrimType -> Int
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
et
  ppr (Replicate Shape
ne SubExp
ve) =
    String -> Doc
text String
"replicate" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [Shape -> Doc
forall a. Pretty a => a -> Doc
ppr Shape
ne, Doc -> Doc
align (SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
ve)]
  ppr (Scratch PrimType
t [SubExp]
shape) =
    String -> Doc
text String
"scratch" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply (PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t Doc -> [Doc] -> [Doc]
forall a. a -> [a] -> [a]
: (SubExp -> Doc) -> [SubExp] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr [SubExp]
shape)
  ppr (Reshape ShapeChange SubExp
shape VName
e) =
    String -> Doc
text String
"reshape" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [[Doc] -> Doc
apply ((DimChange SubExp -> Doc) -> ShapeChange SubExp -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map DimChange SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr ShapeChange SubExp
shape), VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
e]
  ppr (Rearrange [Int]
perm VName
e) =
    String -> Doc
text String
"rearrange" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [[Doc] -> Doc
apply ((Int -> Doc) -> [Int] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Doc
forall a. Pretty a => a -> Doc
ppr [Int]
perm), VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
e]
  ppr (Rotate [SubExp]
es VName
e) =
    String -> Doc
text String
"rotate" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [[Doc] -> Doc
apply ((SubExp -> Doc) -> [SubExp] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr [SubExp]
es), VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
e]
  ppr (Concat Int
i VName
x [VName]
ys SubExp
w) =
    String -> Doc
text String
"concat" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
"@" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Int -> Doc
forall a. Pretty a => a -> Doc
ppr Int
i Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply (SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
w Doc -> [Doc] -> [Doc]
forall a. a -> [a] -> [a]
: VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
x Doc -> [Doc] -> [Doc]
forall a. a -> [a] -> [a]
: (VName -> Doc) -> [VName] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map VName -> Doc
forall a. Pretty a => a -> Doc
ppr [VName]
ys)
  ppr (Copy VName
e) = String -> Doc
text String
"copy" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens (VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
e)
  ppr (Manifest [Int]
perm VName
e) = String -> Doc
text String
"manifest" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [[Doc] -> Doc
apply ((Int -> Doc) -> [Int] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Doc
forall a. Pretty a => a -> Doc
ppr [Int]
perm), VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
e]
  ppr (Assert SubExp
e ErrorMsg SubExp
msg (SrcLoc
loc, [SrcLoc]
_)) =
    String -> Doc
text String
"assert" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
e, ErrorMsg SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr ErrorMsg SubExp
msg, String -> Doc
text (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ String -> String
forall a. Show a => a -> String
show (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc]
  ppr (UpdateAcc VName
acc [SubExp]
is [SubExp]
v) =
    String -> Doc
text String
"update_acc" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
acc, [SubExp] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [SubExp]
is, [SubExp] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [SubExp]
v]

instance Pretty a => Pretty (ErrorMsg a) where
  ppr :: ErrorMsg a -> Doc
ppr (ErrorMsg [ErrorMsgPart a]
parts) = Doc -> Doc
braces (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ Doc -> Doc
align (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (ErrorMsgPart a -> Doc) -> [ErrorMsgPart a] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map ErrorMsgPart a -> Doc
forall a. Pretty a => ErrorMsgPart a -> Doc
p [ErrorMsgPart a]
parts
    where
      p :: ErrorMsgPart a -> Doc
p (ErrorString String
s) = String -> Doc
text (String -> Doc) -> String -> Doc
forall a b. (a -> b) -> a -> b
$ String -> String
forall a. Show a => a -> String
show String
s
      p (ErrorVal PrimType
t a
x) = a -> Doc
forall a. Pretty a => a -> Doc
ppr a
x Doc -> Doc -> Doc
<+> Doc
colon Doc -> Doc -> Doc
<+> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t

instance PrettyRep rep => Pretty (Exp rep) where
  ppr :: Exp rep -> Doc
ppr (If SubExp
c BodyT rep
t BodyT rep
f (IfDec [BranchType rep]
ret IfSort
ifsort)) =
    String -> Doc
text String
"if" Doc -> Doc -> Doc
<+> Doc
info' Doc -> Doc -> Doc
<+> SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
c
      Doc -> Doc -> Doc
</> String -> Doc
text String
"then"
      Doc -> Doc -> Doc
<+> BodyT rep -> Doc
forall rep. PrettyRep rep => Body rep -> Doc
maybeNest BodyT rep
t
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"else"
      Doc -> Doc -> Doc
<+> BodyT rep -> Doc
forall rep. PrettyRep rep => Body rep -> Doc
maybeNest BodyT rep
f
      Doc -> Doc -> Doc
</> Doc
colon
      Doc -> Doc -> Doc
<+> [BranchType rep] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [BranchType rep]
ret
    where
      info' :: Doc
info' = case IfSort
ifsort of
        IfSort
IfNormal -> Doc
forall a. Monoid a => a
mempty
        IfSort
IfFallback -> String -> Doc
text String
"<fallback>"
        IfSort
IfEquiv -> String -> Doc
text String
"<equiv>"
      maybeNest :: BodyT rep -> Doc
maybeNest BodyT rep
b
        | Seq (Stm rep) -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Seq (Stm rep) -> Bool) -> Seq (Stm rep) -> Bool
forall a b. (a -> b) -> a -> b
$ BodyT rep -> Seq (Stm rep)
forall rep. BodyT rep -> Stms rep
bodyStms BodyT rep
b = BodyT rep -> Doc
forall a. Pretty a => a -> Doc
ppr BodyT rep
b
        | Bool
otherwise = String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ BodyT rep -> Doc
forall a. Pretty a => a -> Doc
ppr BodyT rep
b
  ppr (BasicOp BasicOp
op) = BasicOp -> Doc
forall a. Pretty a => a -> Doc
ppr BasicOp
op
  ppr (Apply Name
fname [(SubExp, Diet)]
args [RetType rep]
ret (Safety
safety, SrcLoc
_, [SrcLoc]
_)) =
    Doc
applykw
      Doc -> Doc -> Doc
<+> String -> Doc
text (Name -> String
nameToString Name
fname)
      Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply (((SubExp, Diet) -> Doc) -> [(SubExp, Diet)] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
align (Doc -> Doc) -> ((SubExp, Diet) -> Doc) -> (SubExp, Diet) -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SubExp, Diet) -> Doc
forall a. Pretty a => (a, Diet) -> Doc
pprArg) [(SubExp, Diet)]
args)
      Doc -> Doc -> Doc
</> Doc
colon
      Doc -> Doc -> Doc
<+> Doc -> Doc
braces ([Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (RetType rep -> Doc) -> [RetType rep] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map RetType rep -> Doc
forall a. Pretty a => a -> Doc
ppr [RetType rep]
ret)
    where
      pprArg :: (a, Diet) -> Doc
pprArg (a
arg, Diet
Consume) = String -> Doc
text String
"*" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> a -> Doc
forall a. Pretty a => a -> Doc
ppr a
arg
      pprArg (a
arg, Diet
_) = a -> Doc
forall a. Pretty a => a -> Doc
ppr a
arg
      applykw :: Doc
applykw = case Safety
safety of
        Safety
Unsafe -> String -> Doc
text String
"apply <unsafe>"
        Safety
Safe -> String -> Doc
text String
"apply"
  ppr (Op Op rep
op) = Op rep -> Doc
forall a. Pretty a => a -> Doc
ppr Op rep
op
  ppr (DoLoop [(FParam rep, SubExp)]
merge LoopForm rep
form BodyT rep
loopbody) =
    String -> Doc
text String
"loop" Doc -> Doc -> Doc
<+> Doc -> Doc
braces ([Doc] -> Doc
commastack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (FParam rep -> Doc) -> [FParam rep] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map FParam rep -> Doc
forall a. Pretty a => a -> Doc
ppr [FParam rep]
params)
      Doc -> Doc -> Doc
<+> Doc
equals
      Doc -> Doc -> Doc
<+> [SubExp] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [SubExp]
args
      Doc -> Doc -> Doc
</> ( case LoopForm rep
form of
              ForLoop VName
i IntType
it SubExp
bound [] ->
                String -> Doc
text String
"for"
                  Doc -> Doc -> Doc
<+> Doc -> Doc
align
                    ( VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
i Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
":" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> IntType -> Doc
forall a. Pretty a => a -> Doc
ppr IntType
it
                        Doc -> Doc -> Doc
<+> String -> Doc
text String
"<"
                        Doc -> Doc -> Doc
<+> Doc -> Doc
align (SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
bound)
                    )
              ForLoop VName
i IntType
it SubExp
bound [(LParam rep, VName)]
loop_vars ->
                String -> Doc
text String
"for"
                  Doc -> Doc -> Doc
<+> Doc -> Doc
align
                    ( VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
i Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> String -> Doc
text String
":" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> IntType -> Doc
forall a. Pretty a => a -> Doc
ppr IntType
it
                        Doc -> Doc -> Doc
<+> String -> Doc
text String
"<"
                        Doc -> Doc -> Doc
<+> Doc -> Doc
align (SubExp -> Doc
forall a. Pretty a => a -> Doc
ppr SubExp
bound)
                        Doc -> Doc -> Doc
</> [Doc] -> Doc
stack (((LParam rep, VName) -> Doc) -> [(LParam rep, VName)] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (LParam rep, VName) -> Doc
forall a a. (Pretty a, Pretty a) => (a, a) -> Doc
pprLoopVar [(LParam rep, VName)]
loop_vars)
                    )
              WhileLoop VName
cond ->
                String -> Doc
text String
"while" Doc -> Doc -> Doc
<+> VName -> Doc
forall a. Pretty a => a -> Doc
ppr VName
cond
          )
      Doc -> Doc -> Doc
<+> String -> Doc
text String
"do"
      Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (BodyT rep -> Doc
forall a. Pretty a => a -> Doc
ppr BodyT rep
loopbody)
    where
      ([FParam rep]
params, [SubExp]
args) = [(FParam rep, SubExp)] -> ([FParam rep], [SubExp])
forall a b. [(a, b)] -> ([a], [b])
unzip [(FParam rep, SubExp)]
merge
      pprLoopVar :: (a, a) -> Doc
pprLoopVar (a
p, a
a) = a -> Doc
forall a. Pretty a => a -> Doc
ppr a
p Doc -> Doc -> Doc
<+> String -> Doc
text String
"in" Doc -> Doc -> Doc
<+> a -> Doc
forall a. Pretty a => a -> Doc
ppr a
a
  ppr (WithAcc [WithAccInput rep]
inputs Lambda rep
lam) =
    String -> Doc
text String
"with_acc"
      Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens (Doc -> Doc
braces ([Doc] -> Doc
commastack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (WithAccInput rep -> Doc) -> [WithAccInput rep] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map WithAccInput rep -> Doc
forall a a a a.
(Pretty a, Pretty a, Pretty a, Pretty a) =>
(a, [a], Maybe (a, [a])) -> Doc
ppInput [WithAccInput rep]
inputs) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma Doc -> Doc -> Doc
</> Lambda rep -> Doc
forall a. Pretty a => a -> Doc
ppr Lambda rep
lam)
    where
      ppInput :: (a, [a], Maybe (a, [a])) -> Doc
ppInput (a
shape, [a]
arrs, Maybe (a, [a])
op) =
        Doc -> Doc
parens
          ( a -> Doc
forall a. Pretty a => a -> Doc
ppr a
shape Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma Doc -> Doc -> Doc
<+> [a] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [a]
arrs
              Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> case Maybe (a, [a])
op of
                Maybe (a, [a])
Nothing -> Doc
forall a. Monoid a => a
mempty
                Just (a
op', [a]
nes) ->
                  Doc
comma Doc -> Doc -> Doc
</> Doc -> Doc
parens (a -> Doc
forall a. Pretty a => a -> Doc
ppr a
op' Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma Doc -> Doc -> Doc
</> [Doc] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' ((a -> Doc) -> [a] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map a -> Doc
forall a. Pretty a => a -> Doc
ppr [a]
nes))
          )

instance PrettyRep rep => Pretty (Lambda rep) where
  ppr :: Lambda rep -> Doc
ppr (Lambda [] (Body BodyDec rep
_ Stms rep
stms []) []) | Stms rep
stms Stms rep -> Stms rep -> Bool
forall a. Eq a => a -> a -> Bool
== Stms rep
forall a. Monoid a => a
mempty = String -> Doc
text String
"nilFn"
  ppr (Lambda [LParam rep]
params BodyT rep
body [Type]
rettype) =
    String -> Doc
text String
"\\" Doc -> Doc -> Doc
<+> [LParam rep] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [LParam rep]
params
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (Doc
colon Doc -> Doc -> Doc
<+> [Type] -> Doc
forall a. Pretty a => [a] -> Doc
ppTupleLines' [Type]
rettype Doc -> Doc -> Doc
<+> String -> Doc
text String
"->")
      Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (BodyT rep -> Doc
forall a. Pretty a => a -> Doc
ppr BodyT rep
body)

instance Pretty EntryPointType where
  ppr :: EntryPointType -> Doc
ppr (TypeDirect Uniqueness
u) = Uniqueness -> Doc
forall a. Pretty a => a -> Doc
ppr Uniqueness
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"direct"
  ppr (TypeUnsigned Uniqueness
u) = Uniqueness -> Doc
forall a. Pretty a => a -> Doc
ppr Uniqueness
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"unsigned"
  ppr (TypeOpaque Uniqueness
u String
desc Int
n) = Uniqueness -> Doc
forall a. Pretty a => a -> Doc
ppr Uniqueness
u Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"opaque" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> [Doc] -> Doc
apply [String -> Doc
forall a. Pretty a => a -> Doc
ppr (String -> String
forall a. Show a => a -> String
show String
desc), Int -> Doc
forall a. Pretty a => a -> Doc
ppr Int
n]

instance Pretty EntryParam where
  ppr :: EntryParam -> Doc
ppr (EntryParam Name
name EntryPointType
t) = Name -> Doc
forall a. Pretty a => a -> Doc
ppr Name
name Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
colon Doc -> Doc -> Doc
<+> EntryPointType -> Doc
forall a. Pretty a => a -> Doc
ppr EntryPointType
t

instance PrettyRep rep => Pretty (FunDef rep) where
  ppr :: FunDef rep -> Doc
ppr (FunDef Maybe EntryPoint
entry Attrs
attrs Name
name [RetType rep]
rettype [FParam rep]
fparams BodyT rep
body) =
    [Doc] -> Doc -> Doc
annot (Attrs -> [Doc]
attrAnnots Attrs
attrs) (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$
      Doc
fun
        Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (String -> Doc
text (Name -> String
nameToString Name
name))
        Doc -> Doc -> Doc
<+> [Doc] -> Doc
apply ((FParam rep -> Doc) -> [FParam rep] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map FParam rep -> Doc
forall a. Pretty a => a -> Doc
ppr [FParam rep]
fparams)
        Doc -> Doc -> Doc
</> Int -> Doc -> Doc
indent Int
2 (Doc
colon Doc -> Doc -> Doc
<+> Doc -> Doc
align ([RetType rep] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [RetType rep]
rettype))
        Doc -> Doc -> Doc
<+> Doc
equals
        Doc -> Doc -> Doc
<+> String -> String -> Doc -> Doc
nestedBlock String
"{" String
"}" (BodyT rep -> Doc
forall a. Pretty a => a -> Doc
ppr BodyT rep
body)
    where
      fun :: Doc
fun = case Maybe EntryPoint
entry of
        Maybe EntryPoint
Nothing -> Doc
"fun"
        Just (Name
p_name, [EntryParam]
p_entry, [EntryPointType]
ret_entry) ->
          Doc
"entry"
            Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens
              ( Doc
"\"" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Name -> Doc
forall a. Pretty a => a -> Doc
ppr Name
p_name Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"\"" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma
                  Doc -> Doc -> Doc
</> [EntryParam] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [EntryParam]
p_entry Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
comma
                  Doc -> Doc -> Doc
</> [EntryPointType] -> Doc
forall a. Pretty a => [a] -> Doc
ppTuple' [EntryPointType]
ret_entry
              )

instance PrettyRep rep => Pretty (Prog rep) where
  ppr :: Prog rep -> Doc
ppr (Prog Stms rep
consts [FunDef rep]
funs) =
    [Doc] -> Doc
stack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ Doc -> [Doc] -> [Doc]
punctuate Doc
line ([Doc] -> [Doc]) -> [Doc] -> [Doc]
forall a b. (a -> b) -> a -> b
$ Stms rep -> Doc
forall a. Pretty a => a -> Doc
ppr Stms rep
consts Doc -> [Doc] -> [Doc]
forall a. a -> [a] -> [a]
: (FunDef rep -> Doc) -> [FunDef rep] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map FunDef rep -> Doc
forall a. Pretty a => a -> Doc
ppr [FunDef rep]
funs

instance Pretty d => Pretty (DimChange d) where
  ppr :: DimChange d -> Doc
ppr (DimCoercion d
se) = String -> Doc
text String
"~" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> d -> Doc
forall a. Pretty a => a -> Doc
ppr d
se
  ppr (DimNew d
se) = d -> Doc
forall a. Pretty a => a -> Doc
ppr d
se

instance Pretty d => Pretty (DimIndex d) where
  ppr :: DimIndex d -> Doc
ppr (DimFix d
i) = d -> Doc
forall a. Pretty a => a -> Doc
ppr d
i
  ppr (DimSlice d
i d
n d
s) = d -> Doc
forall a. Pretty a => a -> Doc
ppr d
i Doc -> Doc -> Doc
<+> String -> Doc
text String
":+" Doc -> Doc -> Doc
<+> d -> Doc
forall a. Pretty a => a -> Doc
ppr d
n Doc -> Doc -> Doc
<+> String -> Doc
text String
"*" Doc -> Doc -> Doc
<+> d -> Doc
forall a. Pretty a => a -> Doc
ppr d
s

-- | Like 'prettyTuple', but produces a 'Doc'.
ppTuple' :: Pretty a => [a] -> Doc
ppTuple' :: [a] -> Doc
ppTuple' [a]
ets = Doc -> Doc
braces (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
commasep ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (a -> Doc) -> [a] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
align (Doc -> Doc) -> (a -> Doc) -> a -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Doc
forall a. Pretty a => a -> Doc
ppr) [a]
ets

-- | Like 'prettyTupleLines', but produces a 'Doc'.
ppTupleLines' :: Pretty a => [a] -> Doc
ppTupleLines' :: [a] -> Doc
ppTupleLines' [a]
ets = Doc -> Doc
braces (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
stack ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ Doc -> [Doc] -> [Doc]
punctuate Doc
comma ([Doc] -> [Doc]) -> [Doc] -> [Doc]
forall a b. (a -> b) -> a -> b
$ (a -> Doc) -> [a] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
align (Doc -> Doc) -> (a -> Doc) -> a -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Doc
forall a. Pretty a => a -> Doc
ppr) [a]
ets