How do I make this algorithm lazier without repeating myself?

Question

(Inspired by my answer to this question.)

Consider this code (it's supposed to find the largest element that's less than or equal to a given input):

data TreeMap v = Leaf | Node Integer v (TreeMap v) (TreeMap v) deriving (Show, Read, Eq, Ord)

closestLess :: Integer -> TreeMap v -> Maybe (Integer, v)
closestLess i = precise Nothing where
  precise :: Maybe (Integer, v) -> TreeMap v -> Maybe (Integer, v)
  precise closestSoFar Leaf = closestSoFar
  precise closestSoFar (Node k v l r) = case i `compare` k of
    LT -> precise closestSoFar l
    EQ -> Just (k, v)
    GT -> precise (Just (k, v)) r

This isn't very lazy. Once the GT case is entered, we know for sure that the final return value will be Just something rather than Nothing, but the Just still isn't available until the end. I'd like to make this lazier so that the Just is available as soon as the GT case is entered. My test case for this is that I want Data.Maybe.isJust $ closestLess 5 (Node 3 () Leaf undefined) to evaluate to True rather than bottoming. Here's one way I can think to do this:

data TreeMap v = Leaf | Node Integer v (TreeMap v) (TreeMap v) deriving (Show, Read, Eq, Ord)

closestLess :: Integer -> TreeMap v -> Maybe (Integer, v)
closestLess _ Leaf = Nothing
closestLess i (Node k v l r) = case i `compare` k of
  LT -> closestLess i l
  EQ -> Just (k, v)
  GT -> Just (precise (k, v) r)
  where
    precise :: (Integer, v) -> TreeMap v -> (Integer, v)
    precise closestSoFar Leaf = closestSoFar
    precise closestSoFar (Node k v l r) = case i `compare` k of
      LT -> precise closestSoFar l
      EQ -> (k, v)
      GT -> precise (k, v) r

However, I'm now repeating myself: the core logic is now in both closestLess and in precise. How can I write this so that it's lazy but without repeating myself?

score 4 · Accepted Answer · answered Dec 14 '19 at 22:23

Rather than use explicit wrappers, you can leverage the type system. Note that the version of precise that uses Maybe for your first code snippet:

precise :: Maybe (Integer, v) -> TreeMap v -> Maybe (Integer, v)
precise closestSoFar Leaf = closestSoFar
precise closestSoFar (Node k v l r) = case i `compare` k of
  LT -> precise closestSoFar l
  EQ -> Just (k, v)
  GT -> precise (Just (k, v)) r

is almost exactly the same algorithm as the version of precise without Maybe from your second code snippet, which could be written in the Identity functor as:

precise :: Identity (Integer, v) -> TreeMap v -> Identity (Integer, v)
precise closestSoFar Leaf = closestSoFar
precise closestSoFar (Node k v l r) = case i `compare` k of
  LT -> precise closestSoFar l
  EQ -> Identity (k, v)
  GT -> precise (Identity (k, v)) r

These can be unified into a version polymorphic in the Applicative:

precise :: (Applicative f) => f (Integer, v) -> TreeMap v -> f (Integer, v)
precise closestSoFar Leaf = closestSoFar
precise closestSoFar (Node k v l r) = case i `compare` k of
  LT -> precise closestSoFar l
  EQ -> pure (k, v)
  GT -> precise (pure (k, v)) r

By itself, that doesn't accomplish much, but if we know that the GT branch will always return a value, we can force it to run in the Identity functor, regardless of the starting functor. That is, we can start in the Maybe functor but recurse into the Identity functor in the GT branch:

closestLess :: Integer -> TreeMap v -> Maybe (Integer, v)
closestLess i = precise Nothing
  where
    precise :: (Applicative t) => t (Integer, v) -> TreeMap v -> t (Integer, v)
    precise closestSoFar Leaf = closestSoFar
    precise closestSoFar (Node k v l r) = case i `compare` k of
      LT -> precise closestSoFar l
      EQ -> pure (k, v)
      GT -> pure . runIdentity $ precise (Identity (k, v)) r

This works fine with your test case:

> isJust $ closestLess 5 (Node 3 () Leaf undefined)
True

and is a nice example of polymorphic recursion.

Another nice thing about this approach from a performance point of view is that the -ddump-simpl shows that there are no wrappers or dictionaries. It's all been erased at the type level with specialized functions for the two functors:

closestLess
  = \ @ v i eta ->
      letrec {
        $sprecise
        $sprecise
          = \ @ v1 closestSoFar ds ->
              case ds of {
                Leaf -> closestSoFar;
                Node k v2 l r ->
                  case compareInteger i k of {
                    LT -> $sprecise closestSoFar l;
                    EQ -> (k, v2) `cast` <Co:5>;
                    GT -> $sprecise ((k, v2) `cast` <Co:5>) r
                  }
              }; } in
      letrec {
        $sprecise1
        $sprecise1
          = \ @ v1 closestSoFar ds ->
              case ds of {
                Leaf -> closestSoFar;
                Node k v2 l r ->
                  case compareInteger i k of {
                    LT -> $sprecise1 closestSoFar l;
                    EQ -> Just (k, v2);
                    GT -> Just (($sprecise ((k, v2) `cast` <Co:5>) r) `cast` <Co:4>)
                  }
              }; } in
      $sprecise1 Nothing eta

score 3 · Answer 2 · answered Dec 14 '19 at 20:17

Starting from my non-lazy implementation, I first refactored precise to receive Just as an argument, and generalized its type accordingly:

data TreeMap v = Leaf | Node Integer v (TreeMap v) (TreeMap v) deriving (Show, Read, Eq, Ord)

closestLess :: Integer -> TreeMap v -> Maybe (Integer, v)
closestLess i = precise Just Nothing where
  precise :: ((Integer, v) -> t) -> t -> TreeMap v -> t
  precise _ closestSoFar Leaf = closestSoFar
  precise wrap closestSoFar (Node k v l r) = case i `compare` k of
    LT -> precise wrap closestSoFar l
    EQ -> wrap (k, v)
    GT -> precise wrap (wrap (k, v)) r

Then, I changed it to do wrap early and call itself with id in the GT case:

data TreeMap v = Leaf | Node Integer v (TreeMap v) (TreeMap v) deriving (Show, Read, Eq, Ord)

closestLess :: Integer -> TreeMap v -> Maybe (Integer, v)
closestLess i = precise Just Nothing where
  precise :: ((Integer, v) -> t) -> t -> TreeMap v -> t
  precise _ closestSoFar Leaf = closestSoFar
  precise wrap closestSoFar (Node k v l r) = case i `compare` k of
    LT -> precise wrap closestSoFar l
    EQ -> wrap (k, v)
    GT -> wrap (precise id (k, v) r)

This still works exactly as before, except for the benefit of the added laziness.

Are all those `id`s in the middle between `Just` and the final `(k,v)` eliminated by the compiler? probably not, functions are supposed to be opaque, and you could've (type-feasibly) used `first (1+)` instead of `id` for all the compiler knows. but it makes for a *compact* code... of course, my code is the unraveling and specification of yours here, with the additional simplification (the elimination of the `id`s). also very interesting how the more general type serves as a constraint, a relation between the values involved (not tight enough though, with `first (1+)` being allowed as `wrap`). — Will Ness, Dec 15 '19 at 19:52
(contd.) your polymorphic `precise` is used at two types, directly corresponding to the two specialized functions used in the more verbose variant. nice interplay there. Also, I wouldn't call this CPS, `wrap` isn't used as a continuation, it's not built up "on the inside", it's stacked - by recursion - on the outside. Maybe if it *were* used as continuation you could get rid of those extraneous `id`s... btw we can see here once again that old pattern of functional argument used as indicator of what to do, switching between the two courses of action (`Just` or `id`). — Will Ness, Dec 15 '19 at 20:09

Li-yao Xia · Answer 3 · 2019-12-15T15:26:23.983

I think the CPS version you answered with yourself is the best but for completeness here are a few more ideas. (EDIT: Buhr's answer is now the most performant.)

The first idea is to get rid of the "closestSoFar" accumulator, and instead let the GT case handle all the logic of picking the rightmost value smallest than the argument. In this form, the GT case can directly return a Just:

closestLess1 :: Integer -> TreeMap v -> Maybe (Integer, v)
closestLess1 _ Leaf = Nothing
closestLess1 i (Node k v l r) =
  case i `compare` k of
    LT -> closestLess1 i l
    EQ -> Just (k, v)
    GT -> Just (fromMaybe (k, v) (closestLess1 i r))

This is simpler, but takes a bit more space on the stack when you hit a lot of GT cases. Technically you could even use that fromMaybe in the accumulator form (i.e., replacing the fromJust implicit in luqui's answer), but that would be a redundant, unreachable branch.

The other idea that there's really two "phases" of the algorithm, one before and one after you hit a GT, so you parameterize it by a boolean to represent these two phases, and use dependent types to encode the invariant that there will always be a result in the second phase.

data SBool (b :: Bool) where
  STrue :: SBool 'True
  SFalse :: SBool 'False

type family MaybeUnless (b :: Bool) a where
  MaybeUnless 'True a = a
  MaybeUnless 'False a = Maybe a

ret :: SBool b -> a -> MaybeUnless b a
ret SFalse = Just
ret STrue = id

closestLess2 :: Integer -> TreeMap v -> Maybe (Integer, v)
closestLess2 i = precise SFalse Nothing where
  precise :: SBool b -> MaybeUnless b (Integer, v) -> TreeMap v -> MaybeUnless b (Integer, v)
  precise _ closestSoFar Leaf = closestSoFar
  precise b closestSoFar (Node k v l r) = case i `compare` k of
    LT -> precise b closestSoFar l
    EQ -> ret b (k, v)
    GT -> ret b (precise STrue (k, v) r)

score 2 · Answer 4 · answered Dec 14 '19 at 19:00

2

How about

GT -> let Just v = precise (Just (k,v) r) in Just v

?

answered Dec 14 '19 at 19:00

luqui

59,485
12
145
204

Because that's an incomplete pattern match. Even if my function is a whole is total, I don't like pieces of it being partial. – Joseph Sible-Reinstate Monica Dec 14 '19 at 19:04
So you said "we know for sure" still with some doubt. Perhaps that is healthy. – luqui Dec 14 '19 at 19:09
We do know for sure, given that my second code block in my question always returns `Just` yet is total. I know that your solution as written is in fact total, but it's brittle in that a seemingly-safe modification could then result in it bottoming. – Joseph Sible-Reinstate Monica Dec 14 '19 at 19:10
This will also slightly slow down the program, as GHC can't prove that it will always be `Just`, so it will add a test to make sure it's not `Nothing` every time it recurses through. – Joseph Sible-Reinstate Monica Dec 14 '19 at 22:51

score 1 · Answer 5 · answered Dec 15 '19 at 11:30

Not only do we always know Just, after its first discovery, we also always know Nothing until then. That's actually two different "logics".

So, we go left first of all, so make that explicit:

data TreeMap v = Leaf | Node Integer v (TreeMap v) (TreeMap v) 
                 deriving (Show, Read, Eq, Ord)

closestLess :: Integer 
            -> TreeMap v 
            -> Maybe (Integer, v)
closestLess i = goLeft 
  where
  goLeft :: TreeMap v -> Maybe (Integer, v)
  goLeft n@(Node k v l _) = case i `compare` k of
          LT -> goLeft l
          _  -> Just (precise (k, v) n)
  goLeft Leaf = Nothing

  -- no more maybe if we're here
  precise :: (Integer, v) -> TreeMap v -> (Integer, v)
  precise closestSoFar Leaf           = closestSoFar
  precise closestSoFar (Node k v l r) = case i `compare` k of
        LT -> precise closestSoFar l
        EQ -> (k, v)
        GT -> precise (k, v) r

The price is we repeat at most one step at most once.

How do I make this algorithm lazier without repeating myself?

5 Answers5

Linked