6 | module FS.Chunk

  8 | import FS.Core as P

  9 | import FS.Pull as P

 10 | import Data.List

 11 | import Data.Maybe

 12 | import Data.Nat

 14 | %default total

 27 | public export

 28 | record ChunkSize where

 29 |   [noHints]

 30 |   constructor CS

 31 |   size        : Nat

 32 |   {auto 0 prf : IsSucc size}

 34 | export %inline

 35 | Eq ChunkSize where

 36 |   CS x == CS y = x == y

 38 | export %inline

 39 | Ord ChunkSize where

 40 |   compare (CS x) (CS y) = compare x y

 42 | export %inline

 43 | Show ChunkSize where

 44 |   show = show . size

 47 | fromInteger : (n : Integer) -> ChunkSize

 48 | fromInteger n =

 49 |   case cast {to = Nat} n of

 50 |     k@(S _) => CS k

 51 |     0       => CS 1

 53 | public export %inline %hint

 54 | defaultChunkSize : ChunkSize

 55 | defaultChunkSize = 128

 61 | public export

 62 | data SplitRes : Type -> Type where

 63 |   Middle : (pre, post : c) -> SplitRes c

 64 |   All    : Nat -> SplitRes c

 70 | public export

 71 | interface Monoid c => Chunk (0 c,o : Type) | c where

 72 |   unfoldChunk    : ChunkSize => (s ->  Either r (o,s)) -> s -> UnfoldRes r s c

 73 |   replicateChunk : ChunkSize => o -> c

 74 |   isEmpty        : c -> Bool

 75 |   unconsChunk    : c -> Maybe (o, Maybe c)

 76 |   splitChunkAt   : Nat -> c -> SplitRes c

 77 |   breakChunk     : BreakInstruction -> (o -> Bool) -> c -> BreakRes c

 78 |   filterChunk    : (o -> Bool) -> c -> Maybe c

 81 | nonEmpty : Chunk c o => c -> Maybe c

 82 | nonEmpty v = if isEmpty v then Nothing else Just v

 89 | export %inline

 90 | unfold : ChunkSize => Chunk c o => s -> (s -> Either r (o,s)) -> Pull f c es r

 91 | unfold init g = P.unfold init (unfoldChunk g)

 95 | fill : ChunkSize => (0 c : _) -> Chunk c o => o -> Pull f c es ()

 96 | fill _ v = P.fill (replicateChunk v)

100 | iterate : ChunkSize => (0 c : _) -> Chunk c o => o -> (o -> o) -> Pull f c es ()

101 | iterate _ v f = unfold v (\x => Right (x, f x))

105 | replicate : ChunkSize => (0 c : _) -> Chunk c o => Nat -> o -> Stream f es c

106 | replicate _ n v =

107 |   Chunk.unfold n $ \case

108 |     0   => Left ()

109 |     S k => Right (v,k)

115 | parameters {auto chnk : Chunk c o}

117 |   ||| Splits the very first element from the head of a `Pull`

118 |   export

119 |   uncons : Pull f c es r -> Pull f q es (Either r (o, Pull f c es r))

120 |   uncons p =

121 |     assert_total $ P.uncons p >>= \case

122 |       Left x => pure (Left x)

123 |       Right (vs,q) => case unconsChunk vs of

124 |         Just (el,rem) => pure $ Right (el,consMaybe rem q)

125 |         Nothing => Chunk.uncons q

127 |   ||| Emits the first `n` elements of a `Pull`, returning the remainder.

128 |   export

129 |   splitAt : Nat -> Pull f c es r -> Pull f c es (Pull f c es r)

130 |   splitAt 0 p = pure p

131 |   splitAt k p =

132 |     assert_total $ P.uncons p >>= \case

133 |       Left v      => pure (pure v)

134 |       Right (vs,q) => case splitChunkAt k vs of

135 |         Middle pre post => cons pre (pure $ cons post q)

136 |         All n           => cons vs (Chunk.splitAt n q)

138 |   ||| Emits the first `n` elements of a `Pull`, returning the remainder.

139 |   export %inline

140 |   take : Nat -> Pull f c es r -> Pull f c es ()

141 |   take 0 = const $ pure ()

142 |   take n = ignore . newScope . Chunk.splitAt n

144 |   ||| Like `take` but limits the number of emitted values.

147 |   export %inline

148 |   limit : Has e es => Lazy e -> Nat -> Pull f c es r -> Pull f c es r

149 |   limit err n p = do

150 |     q      <- Chunk.splitAt n p

151 |     Left v <- peekRes q | Right _ => throw err

152 |     pure v

154 |   ||| Drops the first `n` elements of a `Pull`, returning the

156 |   export %inline

157 |   drop : Nat -> Pull f c es r -> Pull f c es r

158 |   drop k p = join $ drain (Chunk.splitAt k p)

160 |   ||| Emits the first element of a `Pull`, returning the remainder.

161 |   export %inline

162 |   head : Pull f c es r -> Pull f c es ()

163 |   head = Chunk.take 1

165 |   ||| Discards the first element of a `Pull`.

166 |   export %inline

167 |   tail : Pull f c es r -> Pull f c es r

168 |   tail = Chunk.drop 1

170 |   ||| Breaks a pull as soon as the given predicate returns `True` for

179 |   export

180 |   breakFull :

181 |        (orElse : r -> Pull f c es r)

182 |     -> BreakInstruction

183 |     -> (o -> Bool)

184 |     -> Pull f c es r

185 |     -> Pull f c es (Pull f c es r)

186 |   breakFull orElse bi pred = breakPull orElse (breakChunk bi pred)

188 |   ||| Like `breakFull` but fails with an error if the `Pull` is

190 |   export %inline

191 |   forceBreakFull :

192 |        {auto has : Has e es}

193 |     -> Lazy e

194 |     -> BreakInstruction

195 |     -> (o -> Bool)

196 |     -> Pull f c es r

197 |     -> Pull f c es (Pull f c es r)

198 |   forceBreakFull err = breakFull (const $ throw err)

200 |   ||| Emits values until the given predicate returns `True`.

204 |   export

205 |   takeUntil : BreakInstruction -> (o -> Bool) -> Pull f c es r -> Stream f es c

206 |   takeUntil tf pred = ignore . newScope . Chunk.breakFull pure tf pred

208 |   ||| Emits values until the given predicate returns `False`,

210 |   export %inline

211 |   takeWhile : (o -> Bool) -> Pull f c es r -> Stream f es c

212 |   takeWhile pred = Chunk.takeUntil DropHit (not . pred)

214 |   ||| Like `takeWhile` but also includes the first failure.

215 |   export %inline

216 |   takeThrough : (o -> Bool) -> Pull f c es r -> Stream f es c

217 |   takeThrough pred = Chunk.takeUntil TakeHit (not . pred)

219 |   ||| Discards values until the given predicate returns `True`.

223 |   export

224 |   dropUntil : BreakInstruction -> (o -> Bool) -> Pull f c es r -> Pull f c es r

225 |   dropUntil tf pred p = drain (Chunk.breakFull pure tf pred p) >>= id

227 |   ||| Drops values from a stream while the given predicate returns `True`,

229 |   export %inline

230 |   dropWhile : (o -> Bool) -> Pull f c es r -> Pull f c es r

231 |   dropWhile pred = Chunk.dropUntil PostHit (not . pred)

233 |   ||| Like `dropWhile` but also drops the first value where

235 |   export %inline

236 |   dropThrough : (o -> Bool) -> Pull f c es r -> Pull f c es r

237 |   dropThrough pred = Chunk.dropUntil DropHit (not . pred)

239 |   ||| Splits a stream of chunks whenever an element fulfills the given

241 |   export

242 |   split : (o -> Bool) -> Pull f c es r -> Pull f (List c)  es r

243 |   split pred = scanFull neutral impl (map pure . nonEmpty)

244 |     where

245 |       loop : SnocList c -> Maybe c -> (Maybe $ List c, c)

246 |       loop sc Nothing  = (Just $ sc <>> [], neutral)

247 |       loop sc (Just x) =

248 |         assert_total $ case breakChunk DropHit pred x of

249 |           Broken x post => loop (sc :< fromMaybe neutral x) post

250 |           Keep x        => (Just $ sc <>> [], x)

252 |       impl : c -> c -> (Maybe $ List c, c)

253 |       impl pre cur =

254 |         case breakChunk DropHit pred cur of

255 |           Broken x post => loop [<pre <+> fromMaybe neutral x] post

256 |           Keep x        => (Nothing, pre <+> x)

262 | nel : List a -> Maybe (List a)

263 | nel [] = Nothing

264 | nel xs = Just xs

268 | mapMaybe : (o -> Maybe p) -> Pull f (List o) es r -> Pull f (List p) es r

269 | mapMaybe f = P.mapMaybe (nel . mapMaybe f)

272 | export %inline

273 | filter : Chunk c o => (o -> Bool) -> Pull f c es r -> Pull f c es r

274 | filter = P.mapMaybe . filterChunk

277 | export %inline

278 | filterNot : Chunk c o => (o -> Bool) -> Pull f c es r -> Pull f c es r

279 | filterNot pred = Chunk.filter (not . pred)

282 | export %inline

283 | mapOutput : Functor t => (o -> p) -> Pull f (t o) es r -> Pull f (t p) es r

284 | mapOutput = P.mapOutput . map

290 | parameters {auto fld : Foldable t}

292 |   export %inline

293 |   fold : (p -> o -> p) -> p -> Pull f (t o) es r -> Pull f p es r

294 |   fold g = P.fold (foldl g)

296 |   ||| Like `fold` but instead of emitting the result as a single

298 |   export %inline

299 |   foldPair : (p -> o -> p) -> p -> Pull f (t o) es r -> Pull f q es (p,r)

300 |   foldPair g = P.foldPair (foldl g)

302 |   ||| Convenience version of `foldPair` for working on streams.

303 |   export %inline

304 |   foldGet : (p -> o -> p) -> p -> Stream f es (t o) -> Pull f q es p

305 |   foldGet g = P.foldGet (foldl g)

307 |   ||| Like `foldC` but will not emit a value in case of an empty pull.

308 |   export

309 |   fold1 : Chunk (t o) o => (o -> o -> o) -> Pull f (t o) es r -> Pull f o es r

310 |   fold1 g s =

311 |     Chunk.uncons s >>= \case

312 |       Left r      => pure r

313 |       Right (v,q) => Chunk.fold g v q

315 |   ||| Returns `True` if all emitted values of the given stream fulfill

317 |   export %inline

318 |   all : (o -> Bool) -> Pull f (t o) es r -> Stream f es Bool

319 |   all pred = P.all (all pred)

321 |   ||| Returns `True` if any of the emitted values of the given stream fulfills

323 |   export %inline

324 |   any : (o -> Bool) -> Pull f (t o) es r -> Stream f es Bool

325 |   any pred = P.any (any pred)

327 |   ||| Emits the sum over all elements emitted by a `Pull`.

328 |   export %inline

329 |   sum : Num o => Pull f (t o) es r -> Pull f o es r

330 |   sum = Chunk.fold (+) 0

332 |   ||| Emits the number of elements emitted by a `Pull`.

333 |   export %inline

334 |   count : Pull f (t o) es r -> Pull f Nat es r

335 |   count = Chunk.fold (const . S) 0

337 |   ||| Emits the largest output encountered.

338 |   export %inline

339 |   maximum : Chunk (t o) o => Ord o => Pull f (t o) es r -> Pull f o es r

340 |   maximum = Chunk.fold1 max

342 |   ||| Emits the smallest output encountered.

343 |   export %inline

344 |   minimum : Chunk (t o) o => Ord o => Pull f (t o) es r -> Pull f o es r

345 |   minimum = Chunk.fold1 min

347 |   ||| Emits the smallest output encountered.

348 |   export %inline

349 |   mappend : Chunk (t o) o => Semigroup o => Pull f (t o) es r -> Pull f o es r

350 |   mappend = Chunk.fold1 (<+>)

352 |   ||| Accumulates the emitted values over a monoid.

357 |   export %inline

358 |   foldMap : Monoid m => (o -> m) -> Pull f (t o) es r -> Pull f m es r

359 |   foldMap f = Chunk.fold (\v,el => v <+> f el) neutral

365 | public export

366 | interface Functor f => Scan f where

367 |   scanChunk : (s -> o -> (p,s)) -> s -> f o -> (f p, s)

369 | parameters {auto sca : Scan t}

371 |   ||| Computes a stateful running total over all elements emitted by a

373 |   export %inline

374 |   scan : s -> (s -> o -> (p,s)) -> Pull f (t o) es r -> Pull f (t p) es r

375 |   scan ini f = P.scan ini (scanChunk f)

377 |   ||| Zips the input with a running total according to `s`, up to but

380 |   export

381 |   zipWithScan : p -> (p -> o -> p) -> Pull f (t o) es r -> Pull f (t (o,p)) es r

382 |   zipWithScan vp fun =

383 |     Chunk.scan vp $ \vp1,vo => let vp2 := fun vp1 vo in ((vo, vp1),vp2)

385 |   ||| Like `zipWithScan` but the running total is including the current element.

386 |   export

387 |   zipWithScan1 : p -> (p -> o -> p) -> Pull f (t o) es r -> Pull f (t (o,p)) es r

388 |   zipWithScan1 vp fun =

389 |     Chunk.scan vp $ \vp1,vo => let vp2 := fun vp1 vo in ((vo, vp2),vp2)

391 |   ||| Pairs each element in the stream with its 0-based index.

392 |   export %inline

393 |   zipWithIndex : Pull f (t o) es r -> Pull f (t (o,Nat)) es r

394 |   zipWithIndex = Chunk.zipWithScan 0 (\n,_ => S n)

396 |   ||| Pairs each element in the stream with its 1-based count.

397 |   export %inline

398 |   zipWithCount : Pull f (t o) es r -> Pull f (t (o,Nat)) es r

399 |   zipWithCount = Chunk.zipWithScan 1 (\n,_ => S n)

401 |   ||| Emits the count of each element.

402 |   export %inline

403 |   runningCount : Pull f (t o) es r -> Pull f (t Nat) es r

404 |   runningCount = Chunk.scan 1 (\n,_ => (n, S n))

411 | len : SnocList a -> Maybe (List a)

412 | len = nel . (<>> [])

414 | broken : SnocList a -> List a -> BreakRes (List a)

415 | broken sx xs = Broken (len sx) (nel xs)

417 | unfoldList :

418 |      SnocList o

419 |   -> Nat

420 |   -> (s -> Either r (o,s))

421 |   -> s

422 |   -> UnfoldRes r s (List o)

423 | unfoldList sx 0     f x = More (sx <>> []) x

424 | unfoldList sx (S k) f x =

425 |   case f x of

426 |     Right (v,x2) => unfoldList (sx:<v) k f x2

427 |     Left res     => Last res (sx <>> [])

429 | splitAtList : SnocList o -> Nat -> List o -> SplitRes (List o)

430 | splitAtList sx (S k) (h::t) = splitAtList (sx :< h) k t

431 | splitAtList sx n     []     = All n

432 | splitAtList sx 0     xs     = Middle (sx <>> []) xs

434 | breakList :

435 |      SnocList a

436 |   -> BreakInstruction

437 |   -> (a -> Bool)

438 |   -> List a

439 |   -> BreakRes (List a)

440 | breakList sx b f []        = Keep (sx <>> [])

441 | breakList sx b f (x :: xs) =

442 |   case f x of

443 |     True => case b of

444 |       TakeHit => broken (sx :< x) xs

445 |       PostHit => broken sx (x::xs)

446 |       DropHit => broken sx xs

447 |     False => breakList (sx :< x) b f xs

450 | Chunk (List a) a where

451 |   unfoldChunk @{CS (S n)} f x =

452 |     case f x of

453 |       Left res     => Done res

454 |       Right (v,x2) => unfoldList [<v] n f x2

456 |   replicateChunk @{CS n} = List.replicate n

458 |   isEmpty [] = True

459 |   isEmpty _  = False

461 |   unconsChunk []     = Nothing

462 |   unconsChunk (h::t) = Just (h, nel t)

464 |   splitChunkAt = splitAtList [<]

466 |   breakChunk = breakList [<]

468 |   filterChunk pred = nel . filter pred

470 | scanListImpl : SnocList p -> (s -> o -> (p,s)) -> s -> List o -> (List p,s)

471 | scanListImpl sx f v []        = (sx <>> [], v)

472 | scanListImpl sx f v (x :: xs) =

473 |   let (vp,v2) := f v x

474 |   in scanListImpl (sx :< vp) f v2 xs

477 | Scan List where

478 |   scanChunk = scanListImpl [<]

484 | public export

485 | data ZipRes : (a,b,c : Type) -> Type where

486 |   ZB : List c -> ZipRes a b c

487 |   ZL : List a -> List c -> ZipRes a b c

488 |   ZR : List b -> List c -> ZipRes a b c

490 | zipImpl : SnocList c -> (a -> b -> c) -> List a -> List b -> ZipRes a b c

491 | zipImpl sx f (x::xs) (y::ys) = zipImpl (sx :< f x y) f xs ys

492 | zipImpl sx f [] [] = ZB (sx <>> [])

493 | zipImpl sx f xs [] = ZL xs (sx <>> [])

494 | zipImpl sx f [] ys = ZR ys (sx <>> [])

496 | zipWith_ :

497 |      (this : Stream f es (List o1))

498 |   -> (that : Stream f es (List o2))

499 |   -> (k1   : ZipWithLeft f es (List o1) (List o3))

500 |   -> (k2   : ZipWithLeft f es (List o2) (List o3))

501 |   -> (k3   : Stream f es (List o2) -> Stream f es (List o3))

502 |   -> (fun  : o1 -> o2 -> o3)

503 |   -> Stream f es (List o3)

504 | zipWith_ this that k1 k2 k3 fun = do

505 |   Just l1 <- stepLeg this | Nothing => k3 that

506 |   Just l2 <- stepLeg that | Nothing => k1 l1.out l1.pull

507 |   go l1 l2

509 |   where

510 |     go : StepLeg f es (List o1) -> StepLeg f es (List o2) -> Stream f es (List o3)

511 |     go l1 l2 =

512 |       case zipImpl [<] fun l1.out l2.out of

513 |         ZB cs     => do

514 |           emit cs

515 |           Just l12 <- step l1 | Nothing => k3 l2.pull

516 |           Just l22 <- step l2 | Nothing => k1 l12.out l12.pull

517 |           assert_total $ go l12 l22

518 |         ZL as cs => do

519 |           emit cs

520 |           Just l22 <- step l2 | Nothing => k1 as l1.pull

521 |           assert_total $ go ({out := as} l1) l22

522 |         ZR bs cs => do

523 |           emit cs

524 |           Just l12 <- step l1 | Nothing => k2 bs l2.pull

525 |           assert_total $ go l12 ({out := bs} l2)

531 | zipAllWith :

532 |      (pad1 : o1)

533 |   -> (pad2 : o2)

534 |   -> (o1 -> o2 -> o3)

535 |   -> (this : Stream f es (List o1))

536 |   -> (that : Stream f es (List o2))

537 |   -> Stream f es (List o3)

538 | zipAllWith pad1 pad2 fun this that =

539 |  let kL := Chunk.mapOutput (`fun` pad2)

540 |      kR := Chunk.mapOutput (fun pad1)

541 |      k1 := \x,y => emit (map (`fun` pad2) x) >> kL y

542 |      k2 := \x,y => emit (map (fun pad1) x) >> kR y

543 |   in zipWith_ this that k1 k2 kR fun

546 | export %inline

547 | zipAll :

548 |      (pad1 : o1)

549 |   -> (pad2 : o2)

550 |   -> (this : Stream f es (List o1))

551 |   -> (that : Stream f es (List o2))

552 |   -> Stream f es (List (o1,o2))

553 | zipAll pad1 pad2 = zipAllWith pad1 pad2 MkPair

557 | export %inline

558 | zipWith :

559 |      (o1 -> o2 -> o3)

560 |   -> (this : Stream f es (List o1))

561 |   -> (that : Stream f es (List o2))

562 |   -> Stream f es (List o3)

563 | zipWith fun this that =

564 |   zipWith_ this that (\_,_ => pure ()) (\_,_ => pure ()) (\_ => pure ()) fun

567 | export %inline

568 | zip :

569 |      (this : Stream f es (List o1))

570 |   -> (that : Stream f es (List o2))

571 |   -> Stream f es (List (o1,o2))

572 | zip = zipWith MkPair

574 | export %inline

575 | zipRight :

576 |      (this : Stream f es (List o1))

577 |   -> (that : Stream f es (List o2))

578 |   -> Stream f es (List o2)

579 | zipRight = Chunk.zipWith (\_ => id)

581 | export %inline

582 | zipLeft :

583 |      (this : Stream f es (List o1))

584 |   -> (that : Stream f es (List o2))

585 |   -> Stream f es (List o1)

586 | zipLeft = Chunk.zipWith const