nchoosek.hs

{-# language LambdaCase #-}

import Control.Monad.Trans.Class ( lift )
import Control.Monad.Trans.State.Strict
import Data.Functor.Identity
import ListT


-- | Choose as many elements as required to "fill" an applicative functor:
--
-- >>> choose @[] [ 1, 2, 3 ]
-- [[1],[2],[3]]
--
-- >>> choose @Identity [ 1, 2, 3 ]
-- [ Identity 1,Identity 2,Identity 3 ]
--
-- >>> choose @V3 [ 1, 2, 3 ]
-- [ V3 1 2 3,V3 1 3 2,V3 2 1 3,V3 2 3 1,V3 3 1 2,V3 3 2 1 ]

choose :: ( Applicative f, Traversable f ) => [ x ] -> [ f x ]
choose =
  runIdentity
    . toList
    . evalStateT
        ( sequenceA . pure $ do
            ( x, xs ) <-
              lift . fromFoldable =<< picks <$> get

            x <$ put xs
        )


-- | All selections of one element from a list of elements:
--
-- >>> picks [ 1, 2, 3 ]
-- [ ( 1, [ 2, 3 ] ), ( 2, [ 1, 3 ] ), ( 3, [ 1, 2 ] ) ]

picks :: [ a ] -> [ ( a, [ a ] ) ]
picks = \case
  [] ->
    []

  ( x : xs ) ->
    ( x, xs ) : fmap ( fmap ( x : ) ) ( picks xs )