import Data.Set (Set, (\\))
import Data.Set qualified as S

closure :: Ord a => (a -> [a]) -> Set a -> Set a
closure f s = go S.empty s
 where
  go old new | null new  = old
             | otherwise = go old' new'
   where
    old' = old <> new
    new' = S.fromList (concatMap f new) `S.difference` old'

closure' :: Ord a => (a -> [a]) -> Set a -> Set a
closure' f s0 = go (S.toList s0) s0
  where
    go xs s =
      let xs' = concatMap f xs
          s'  = S.union s (S.fromList xs')
      in  if S.size s == S.size s' then s
          else go (S.toList (S.difference s' s)) s'

saturate :: Ord x => [x] -> (x -> S.Set x -> [x]) -> S.Set x
saturate sat0 sat1 = foldl' go S.empty sat0 where
    go xs x
        | x `S.member` xs = xs
        | otherwise = let xs' = S.insert x xs in foldl' go xs' (sat1 x xs')

g = ((:[]) . (`div` 2))

main :: IO ()
main = do
  print $ closure  g (S.singleton 512)
  print $ closure' g (S.singleton 512)
  print $ saturate [512] (\x _ -> g x)