Partition.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DerivingStrategies #-}

import Data.Either

data Decl a
  = TypeSig String a
  | Law String a a
  | FunModel String a
  | TypeModel String a
  | Import a
  | Other a
  deriving stock Show

decls =
  [ TypeSig "Hello" 1,
    Law "Fred" 1 42,
    FunModel "Fun" 2,
    Import 10,
    Other 99,
    Law "Be good" 2 3,
    TypeModel "huh" 0
  ]

partitionThem decls =
  let (laws, others1) = partitionEithers $ extractLaws decls
      (funModels, others2) = partitionEithers $ extractFunModels others1
      (typeModels, leftovers) = partitionEithers $ extractTypeModels others2
   in (laws, funModels, typeModels, leftovers)
  where
    extractLaws = map (\case law@Law {} -> Left law; d -> Right d)
    extractFunModels = map (\case fm@FunModel {} -> Left fm; d -> Right d)
    extractTypeModels = map (\case tm@TypeModel {} -> Left tm; d -> Right d)