newtype State = State
  { f :: Doc -> Doc
  , k :: Int
  , l :: List (Tuple Int DOC)
  -- Based on the generated document, should it keep the result or try a new state
  , next :: Maybe (Tuple (Doc -> Boolean) State)
  }

be :: Int -> Int -> List (Tuple Int DOC) -> Doc
be w k0 l0 = tailRec go $ State { f: id, k: k0, l: l0, next: Nothing }
  where
  go :: State -> Step State Doc
  go (State { f, k, l, next }) = case l of
    List.Nil ->
      -- obtain the result by running the function on the empty document
      let res = f Nil in
      case next of
        -- if this one fails, continue at the next state
        Just (Tuple p st) | not p res ->
          Loop st
        -- otherwise return what we got
        _ -> Done res
    (Cons (Tuple _ NIL) z) -> Loop $
      State { f, k, l: z, next }
    (Cons (Tuple i (APPEND x y)) z) -> Loop $
      State { f, k, l: (Tuple i x) : (Tuple i y) : z, next }
    (Cons (Tuple i (NEST j x)) z) -> Loop $
      State { f, k, l: (Tuple (i + j) x) : z, next }
    (Cons (Tuple _ (TEXT s)) z) -> Loop $
      State { f: f <<< Text s, k: (k + String.length s), l: z, next }
    (Cons (Tuple i LINE) z) -> Loop $
      State { f: f <<< Line i, k: i, l: z, next }
    (Cons (Tuple i (UNION x y)) z) ->
      let
        -- add a continuation that tests whether the previous doc is too long
        trial = Tuple (fits (w - k)) nextSt
        nextSt = State { f, k, l: (Tuple i y) : z, next }
      in  Loop $ State { f, k, l: (Tuple i x) : z, next: Just trial }