Partition.hs

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE LambdaCase #-}
{-# OPTIONS_GHC -Wall #-}

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 :: [Decl Integer]
decls =
  [ TypeSig "Hello" 1,
    Law "Fred" 1 42,
    FunModel "Fun" 2,
    Import 10,
    Other 99,
    Law "Be good" 2 3,
    TypeModel "huh" 0
  ]

-- Loops forever. Every `let` is `letrec` :-(
partitionThem :: [Decl a] -> ([Decl a], [Decl a], [Decl a], [Decl a])
partitionThem ds =
  let (laws, leftovers) = partitionEithers $ extractLaws ds
   in let (funModels, leftovers) = partitionEithers $ extractFunModels leftovers
       in let (typeModels, leftovers) = partitionEithers $ extractTypeModels leftovers
           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)