{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

-- TODO: the calls to 'error' should be replaced with logging/error capabilities.

-- |
-- Module      :  Disco.Pretty
-- Copyright   :  disco team and contributors
-- Maintainer  :  [email protected]
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Various pretty-printing facilities for disco.
module Disco.Pretty (
  module Disco.Pretty.DSL,
  module Disco.Pretty,
  module Disco.Pretty.Prec,
  Doc,
)
where

import Prelude hiding ((<>))

import Data.Bifunctor
import Data.Char (isAlpha)
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NE
import Data.Map (Map)
import qualified Data.Map as M
import Data.Ratio
import Data.Set (Set)
import qualified Data.Set as S
import Disco.Effects.LFresh
import Disco.Pretty.DSL
import Disco.Pretty.Prec
import Disco.Syntax.Operators
import Polysemy
import Polysemy.Reader
import Prettyprinter (Doc)
import Unbound.Generics.LocallyNameless (Name)

------------------------------------------------------------
-- Utilities for handling precedence and associativity

-- | Convenience function combining 'setPA' and 'mparens', since we
--   often want to simultaneously indicate what the precedence and
--   associativity of a term is, and optionally surround it with
--   parentheses depending on the precedence and associativity of its
--   parent.
withPA :: Member (Reader PA) r => PA -> Sem r (Doc ann) -> Sem r (Doc ann)
withPA :: forall (r :: EffectRow) ann.
Member (Reader PA) r =>
PA -> Sem r (Doc ann) -> Sem r (Doc ann)
withPA PA
pa = PA -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (r :: EffectRow) ann.
Member (Reader PA) r =>
PA -> Sem r (Doc ann) -> Sem r (Doc ann)
mparens PA
pa (Sem r (Doc ann) -> Sem r (Doc ann))
-> (Sem r (Doc ann) -> Sem r (Doc ann))
-> Sem r (Doc ann)
-> Sem r (Doc ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PA -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (r :: EffectRow) a.
Member (Reader PA) r =>
PA -> Sem r a -> Sem r a
setPA PA
pa

-- | Locally set the precedence and associativity within a
--   subcomputation.
setPA :: Member (Reader PA) r => PA -> Sem r a -> Sem r a
setPA :: forall (r :: EffectRow) a.
Member (Reader PA) r =>
PA -> Sem r a -> Sem r a
setPA = (PA -> PA) -> Sem r a -> Sem r a
forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local ((PA -> PA) -> Sem r a -> Sem r a)
-> (PA -> PA -> PA) -> PA -> Sem r a -> Sem r a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PA -> PA -> PA
forall a b. a -> b -> a
const

-- | Mark a subcomputation as pretty-printing a term on the left of an
--   operator (so parentheses can be inserted appropriately, depending
--   on the associativity).
lt :: Member (Reader PA) r => Sem r (Doc ann) -> Sem r (Doc ann)
lt :: forall (r :: EffectRow) ann.
Member (Reader PA) r =>
Sem r (Doc ann) -> Sem r (Doc ann)
lt = (PA -> PA) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local (\(PA Int
p BFixity
_) -> Int -> BFixity -> PA
PA Int
p BFixity
InL)

-- | Mark a subcomputation as pretty-printing a term on the right of
--   an operator (so parentheses can be inserted appropriately,
--   depending on the associativity).
rt :: Member (Reader PA) r => Sem r (Doc ann) -> Sem r (Doc ann)
rt :: forall (r :: EffectRow) ann.
Member (Reader PA) r =>
Sem r (Doc ann) -> Sem r (Doc ann)
rt = (PA -> PA) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall i (r :: EffectRow) a.
Member (Reader i) r =>
(i -> i) -> Sem r a -> Sem r a
local (\(PA Int
p BFixity
_) -> Int -> BFixity -> PA
PA Int
p BFixity
InR)

-- | Optionally surround a pretty-printed term with parentheses,
--   depending on its precedence and associativity (given as the 'PA'
--   argument) and that of its context (given by the ambient 'Reader
--   PA' effect).
mparens :: Member (Reader PA) r => PA -> Sem r (Doc ann) -> Sem r (Doc ann)
mparens :: forall (r :: EffectRow) ann.
Member (Reader PA) r =>
PA -> Sem r (Doc ann) -> Sem r (Doc ann)
mparens PA
pa Sem r (Doc ann)
doc = do
  PA
parentPA <- Sem r PA
forall i (r :: EffectRow). Member (Reader i) r => Sem r i
ask
  (if PA
pa PA -> PA -> Bool
`lowerPrec` PA
parentPA then Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
parens else Sem r (Doc ann) -> Sem r (Doc ann)
forall a. a -> a
id) Sem r (Doc ann)
doc

------------------------------------------------------------
-- Pretty type class

class Pretty t where
  pretty :: Members '[Reader PA, LFresh] r => t -> Sem r (Doc ann)

prettyStr :: Pretty t => t -> Sem r String
prettyStr :: forall t (r :: EffectRow). Pretty t => t -> Sem r [Char]
prettyStr = Sem (Reader PA : r) (Doc Any) -> Sem r [Char]
forall (r :: EffectRow) ann.
Sem (Reader PA : r) (Doc ann) -> Sem r [Char]
renderDoc (Sem (Reader PA : r) (Doc Any) -> Sem r [Char])
-> (t -> Sem (Reader PA : r) (Doc Any)) -> t -> Sem r [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sem (LFresh : Reader PA : r) (Doc Any)
-> Sem (Reader PA : r) (Doc Any)
forall (r :: EffectRow) a. Sem (LFresh : r) a -> Sem r a
runLFresh (Sem (LFresh : Reader PA : r) (Doc Any)
 -> Sem (Reader PA : r) (Doc Any))
-> (t -> Sem (LFresh : Reader PA : r) (Doc Any))
-> t
-> Sem (Reader PA : r) (Doc Any)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> Sem (LFresh : Reader PA : r) (Doc Any)
forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
t -> Sem r (Doc ann)
forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty

pretty' :: Pretty t => t -> Sem r (Doc ann)
pretty' :: forall t (r :: EffectRow) ann. Pretty t => t -> Sem r (Doc ann)
pretty' = PA -> Sem (Reader PA : r) (Doc ann) -> Sem r (Doc ann)
forall i (r :: EffectRow) a. i -> Sem (Reader i : r) a -> Sem r a
runReader PA
initPA (Sem (Reader PA : r) (Doc ann) -> Sem r (Doc ann))
-> (t -> Sem (Reader PA : r) (Doc ann)) -> t -> Sem r (Doc ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sem (LFresh : Reader PA : r) (Doc ann)
-> Sem (Reader PA : r) (Doc ann)
forall (r :: EffectRow) a. Sem (LFresh : r) a -> Sem r a
runLFresh (Sem (LFresh : Reader PA : r) (Doc ann)
 -> Sem (Reader PA : r) (Doc ann))
-> (t -> Sem (LFresh : Reader PA : r) (Doc ann))
-> t
-> Sem (Reader PA : r) (Doc ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> Sem (LFresh : Reader PA : r) (Doc ann)
forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
t -> Sem r (Doc ann)
forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty

------------------------------------------------------------
-- Some standard instances

instance Pretty a => Pretty [a] where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
[a] -> Sem r (Doc ann)
pretty = Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
brackets (Sem r (Doc ann) -> Sem r (Doc ann))
-> ([a] -> Sem r (Doc ann)) -> [a] -> Sem r (Doc ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sem r (Doc ann) -> [Sem r (Doc ann)] -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Monad f =>
f (Doc ann) -> [f (Doc ann)] -> f (Doc ann)
intercalate Sem r (Doc ann)
"," ([Sem r (Doc ann)] -> Sem r (Doc ann))
-> ([a] -> [Sem r (Doc ann)]) -> [a] -> Sem r (Doc ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Sem r (Doc ann)) -> [a] -> [Sem r (Doc ann)]
forall a b. (a -> b) -> [a] -> [b]
map a -> Sem r (Doc ann)
forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
a -> Sem r (Doc ann)
forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty

instance Pretty a => Pretty (NonEmpty a) where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
NonEmpty a -> Sem r (Doc ann)
pretty = [a] -> Sem r (Doc ann)
forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
[a] -> Sem r (Doc ann)
forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty ([a] -> Sem r (Doc ann))
-> (NonEmpty a -> [a]) -> NonEmpty a -> Sem r (Doc ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.toList

instance (Pretty k, Pretty v) => Pretty (Map k v) where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
Map k v -> Sem r (Doc ann)
pretty Map k v
m = do
    let es :: [Sem r (Doc ann)]
es = ((k, v) -> Sem r (Doc ann)) -> [(k, v)] -> [Sem r (Doc ann)]
forall a b. (a -> b) -> [a] -> [b]
map (\(k
k, v
v) -> k -> Sem r (Doc ann)
forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
k -> Sem r (Doc ann)
forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty k
k Sem r (Doc ann) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<+> Sem r (Doc ann)
"->" Sem r (Doc ann) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<+> v -> Sem r (Doc ann)
forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
v -> Sem r (Doc ann)
forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty v
v) (Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
M.assocs Map k v
m)
    [Sem r (Doc ann)]
ds <- PA -> Sem r [Sem r (Doc ann)] -> Sem r [Sem r (Doc ann)]
forall (r :: EffectRow) a.
Member (Reader PA) r =>
PA -> Sem r a -> Sem r a
setPA PA
initPA (Sem r [Sem r (Doc ann)] -> Sem r [Sem r (Doc ann)])
-> Sem r [Sem r (Doc ann)] -> Sem r [Sem r (Doc ann)]
forall a b. (a -> b) -> a -> b
$ Sem r (Doc ann) -> [Sem r (Doc ann)] -> Sem r [Sem r (Doc ann)]
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> [f (Doc ann)] -> f [f (Doc ann)]
punctuate Sem r (Doc ann)
"," [Sem r (Doc ann)]
es
    Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
braces ([Sem r (Doc ann)] -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
[f (Doc ann)] -> f (Doc ann)
hsep [Sem r (Doc ann)]
ds)

instance Pretty a => Pretty (Set a) where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
Set a -> Sem r (Doc ann)
pretty = Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann. Functor f => f (Doc ann) -> f (Doc ann)
braces (Sem r (Doc ann) -> Sem r (Doc ann))
-> (Set a -> Sem r (Doc ann)) -> Set a -> Sem r (Doc ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sem r (Doc ann) -> [Sem r (Doc ann)] -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Monad f =>
f (Doc ann) -> [f (Doc ann)] -> f (Doc ann)
intercalate Sem r (Doc ann)
"," ([Sem r (Doc ann)] -> Sem r (Doc ann))
-> (Set a -> [Sem r (Doc ann)]) -> Set a -> Sem r (Doc ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Sem r (Doc ann)) -> [a] -> [Sem r (Doc ann)]
forall a b. (a -> b) -> [a] -> [b]
map a -> Sem r (Doc ann)
forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
a -> Sem r (Doc ann)
forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty ([a] -> [Sem r (Doc ann)])
-> (Set a -> [a]) -> Set a -> [Sem r (Doc ann)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> [a]
forall a. Set a -> [a]
S.toList

------------------------------------------------------------
-- Some Disco instances

instance Pretty (Name a) where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
Name a -> Sem r (Doc ann)
pretty = [Char] -> Sem r (Doc ann)
forall (m :: * -> *) ann. Applicative m => [Char] -> m (Doc ann)
text ([Char] -> Sem r (Doc ann))
-> (Name a -> [Char]) -> Name a -> Sem r (Doc ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name a -> [Char]
forall a. Show a => a -> [Char]
show

instance Pretty TyOp where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
TyOp -> Sem r (Doc ann)
pretty = \case
    TyOp
Enumerate -> [Char] -> Sem r (Doc ann)
forall (m :: * -> *) ann. Applicative m => [Char] -> m (Doc ann)
text [Char]
"enumerate"
    TyOp
Count -> [Char] -> Sem r (Doc ann)
forall (m :: * -> *) ann. Applicative m => [Char] -> m (Doc ann)
text [Char]
"count"

-- | Pretty-print a unary operator, by looking up its concrete syntax
--   in the 'uopMap'.
instance Pretty UOp where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
UOp -> Sem r (Doc ann)
pretty UOp
op = case UOp -> Map UOp OpInfo -> Maybe OpInfo
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup UOp
op Map UOp OpInfo
uopMap of
    Just (OpInfo OpFixity
_ ([Char]
syn : [[Char]]
_) Int
_) ->
      [Char] -> Sem r (Doc ann)
forall (m :: * -> *) ann. Applicative m => [Char] -> m (Doc ann)
text ([Char] -> Sem r (Doc ann)) -> [Char] -> Sem r (Doc ann)
forall a b. (a -> b) -> a -> b
$ [Char]
syn [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (if (Char -> Bool) -> [Char] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
isAlpha [Char]
syn then [Char]
" " else [Char]
"")
    Maybe OpInfo
_ -> [Char] -> Sem r (Doc ann)
forall a. HasCallStack => [Char] -> a
error ([Char] -> Sem r (Doc ann)) -> [Char] -> Sem r (Doc ann)
forall a b. (a -> b) -> a -> b
$ [Char]
"UOp " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ UOp -> [Char]
forall a. Show a => a -> [Char]
show UOp
op [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" not in uopMap!"

-- | Pretty-print a binary operator, by looking up its concrete syntax
--   in the 'bopMap'.
instance Pretty BOp where
  pretty :: forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
BOp -> Sem r (Doc ann)
pretty (Should BOp
op) = BOp -> Sem r (Doc ann)
forall (r :: EffectRow) ann.
Members '[Reader PA, LFresh] r =>
BOp -> Sem r (Doc ann)
forall t (r :: EffectRow) ann.
(Pretty t, Members '[Reader PA, LFresh] r) =>
t -> Sem r (Doc ann)
pretty BOp
op Sem r (Doc ann) -> Sem r (Doc ann) -> Sem r (Doc ann)
forall (f :: * -> *) ann.
Applicative f =>
f (Doc ann) -> f (Doc ann) -> f (Doc ann)
<> Sem r (Doc ann)
"!"
  pretty BOp
op = case BOp -> Map BOp OpInfo -> Maybe OpInfo
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup BOp
op Map BOp OpInfo
bopMap of
    Just (OpInfo OpFixity
_ ([Char]
syn : [[Char]]
_) Int
_) -> [Char] -> Sem r (Doc ann)
forall (m :: * -> *) ann. Applicative m => [Char] -> m (Doc ann)
text [Char]
syn
    Maybe OpInfo
_ -> [Char] -> Sem r (Doc ann)
forall a. HasCallStack => [Char] -> a
error ([Char] -> Sem r (Doc ann)) -> [Char] -> Sem r (Doc ann)
forall a b. (a -> b) -> a -> b
$ [Char]
"BOp " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ BOp -> [Char]
forall a. Show a => a -> [Char]
show BOp
op [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" not in bopMap!"

--------------------------------------------------
-- Pretty-printing rationals + decimals

-- | Pretty-print a rational number as a fraction.
prettyRational :: Rational -> String
prettyRational :: Rational -> [Char]
prettyRational Rational
r
  | Rational -> Integer
forall a. Ratio a -> a
denominator Rational
r Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 = Integer -> [Char]
forall a. Show a => a -> [Char]
show (Rational -> Integer
forall a. Ratio a -> a
numerator Rational
r)
  | Bool
otherwise = Integer -> [Char]
forall a. Show a => a -> [Char]
show (Rational -> Integer
forall a. Ratio a -> a
numerator Rational
r) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"/" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show (Rational -> Integer
forall a. Ratio a -> a
denominator Rational
r)

-- | Pretty-print a rational number using its decimal expansion, in
--   the format @nnn.prefix[rep]...@, with any repeating digits enclosed
--   in square brackets.
prettyDecimal :: Rational -> String
prettyDecimal :: Rational -> [Char]
prettyDecimal Rational
r = [Char]
wholePart [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"." [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
decimalPart
 where
  (Integer
n', Rational
d') = Rational -> (Integer, Rational)
forall b. Integral b => Rational -> (b, Rational)
forall a b. (RealFrac a, Integral b) => a -> (b, a)
properFraction Rational
r :: (Integer, Rational)
  d :: Rational
d = Rational -> Rational
forall a. Num a => a -> a
abs Rational
d'
  n :: Integer
n = Integer -> Integer
forall a. Num a => a -> a
abs Integer
n'
  ([Integer]
expan, Int
len) = Integer -> Integer -> Integer -> ([Integer], Int)
digitalExpansion Integer
10 (Rational -> Integer
forall a. Ratio a -> a
numerator Rational
d) (Rational -> Integer
forall a. Ratio a -> a
denominator Rational
d)
  wholePart :: [Char]
wholePart = (if Rational
d' Rational -> Rational -> Bool
forall a. Ord a => a -> a -> Bool
< Rational
0 then [Char]
"-" else [Char]
"") [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
n
  decimalPart :: [Char]
decimalPart
    | [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
first102 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
101 Bool -> Bool -> Bool
|| [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
first102 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
101 Bool -> Bool -> Bool
&& [Integer] -> Integer
forall a. HasCallStack => [a] -> a
last [Integer]
first102 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
0 =
        (Integer -> [Char]) -> [Integer] -> [Char]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Integer -> [Char]
forall a. Show a => a -> [Char]
show (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
100 [Integer]
expan) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"..."
    | [Integer]
rep [Integer] -> [Integer] -> Bool
forall a. Eq a => a -> a -> Bool
== [Integer
0] =
        if [Integer] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Integer]
pre then [Char]
"0" else (Integer -> [Char]) -> [Integer] -> [Char]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Integer -> [Char]
forall a. Show a => a -> [Char]
show [Integer]
pre
    | Bool
otherwise =
        (Integer -> [Char]) -> [Integer] -> [Char]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Integer -> [Char]
forall a. Show a => a -> [Char]
show [Integer]
pre [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"[" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Integer -> [Char]) -> [Integer] -> [Char]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Integer -> [Char]
forall a. Show a => a -> [Char]
show [Integer]
rep [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"]"
   where
    ([Integer]
pre, [Integer]
rep) = Int -> [Integer] -> ([Integer], [Integer])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
len [Integer]
expan
    first102 :: [Integer]
first102 = Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
102 [Integer]
expan

-- Given a list, find the indices of the list giving the first and
-- second occurrence of the first element to repeat, or Nothing if
-- there are no repeats.
findRep :: Ord a => [a] -> ([a], Int)
findRep :: forall a. Ord a => [a] -> ([a], Int)
findRep = Map a Int -> Int -> [a] -> ([a], Int)
forall a. Ord a => Map a Int -> Int -> [a] -> ([a], Int)
findRep' Map a Int
forall k a. Map k a
M.empty Int
0

findRep' :: Ord a => M.Map a Int -> Int -> [a] -> ([a], Int)
findRep' :: forall a. Ord a => Map a Int -> Int -> [a] -> ([a], Int)
findRep' Map a Int
_ Int
_ [] = [Char] -> ([a], Int)
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible. Empty list in findRep'"
findRep' Map a Int
prevs Int
ix (a
x : [a]
xs)
  | a
x a -> Map a Int -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map a Int
prevs = ([], Map a Int
prevs Map a Int -> a -> Int
forall k a. Ord k => Map k a -> k -> a
M.! a
x)
  | Bool
otherwise = ([a] -> [a]) -> ([a], Int) -> ([a], Int)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (([a], Int) -> ([a], Int)) -> ([a], Int) -> ([a], Int)
forall a b. (a -> b) -> a -> b
$ Map a Int -> Int -> [a] -> ([a], Int)
forall a. Ord a => Map a Int -> Int -> [a] -> ([a], Int)
findRep' (a -> Int -> Map a Int -> Map a Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x Int
ix Map a Int
prevs) (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs

-- | @digitalExpansion b n d@ takes the numerator and denominator of a
--   fraction n/d between 0 and 1, and returns a pair of (1) a list of
--   digits @ds@, and (2) a nonnegative integer k such that @splitAt k
--   ds = (prefix, rep)@, where the infinite base-b expansion of
--   n/d is 0.@(prefix ++ cycle rep)@.  For example,
--
--   > digitalExpansion 10 1 4  = ([2,5,0], 2)
--   > digitalExpansion 10 1 7  = ([1,4,2,8,5,7], 0)
--   > digitalExpansion 10 3 28 = ([1,0,7,1,4,2,8,5], 2)
--   > digitalExpansion 2  1 5  = ([0,0,1,1], 0)
--
--   It works by performing the standard long division algorithm, and
--   looking for the first time that the remainder repeats.
digitalExpansion :: Integer -> Integer -> Integer -> ([Integer], Int)
digitalExpansion :: Integer -> Integer -> Integer -> ([Integer], Int)
digitalExpansion Integer
b Integer
n Integer
d = ([Integer], Int)
digits
 where
  longDivStep :: (Integer, Integer) -> (Integer, Integer)
longDivStep (Integer
_, Integer
r) = (Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
r) Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`divMod` Integer
d
  res :: [(Integer, Integer)]
res = Int -> [(Integer, Integer)] -> [(Integer, Integer)]
forall a. Int -> [a] -> [a]
drop Int
1 ([(Integer, Integer)] -> [(Integer, Integer)])
-> [(Integer, Integer)] -> [(Integer, Integer)]
forall a b. (a -> b) -> a -> b
$ ((Integer, Integer) -> (Integer, Integer))
-> (Integer, Integer) -> [(Integer, Integer)]
forall a. (a -> a) -> a -> [a]
iterate (Integer, Integer) -> (Integer, Integer)
longDivStep (Integer
0, Integer
n)
  digits :: ([Integer], Int)
digits = ([(Integer, Integer)] -> [Integer])
-> ([(Integer, Integer)], Int) -> ([Integer], Int)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (((Integer, Integer) -> Integer)
-> [(Integer, Integer)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst) ([(Integer, Integer)] -> ([(Integer, Integer)], Int)
forall a. Ord a => [a] -> ([a], Int)
findRep [(Integer, Integer)]
res)