Refactor BranchLists.

Now we use Array to store branches. This makes sense because we often
have to do random access (once inference is done). This also vastly
simplifies the awkward BranchList type.

This fixes #10837 and updates submodule utils/haddock.

12 files changed, 117 insertions, 150 deletions

diff --git a/compiler/coreSyn/CoreLint.hs b/compiler/coreSyn/CoreLint.hs
index 2b1118e69b..0b72ff4db2 100644
--- a/compiler/coreSyn/CoreLint.hs
+++ b/compiler/coreSyn/CoreLint.hs
@@ -1336,7 +1336,7 @@ lintCoercion (InstCo co arg_ty)
           _ -> failWithL (ptext (sLit "Bad argument of inst")) }
 
 lintCoercion co@(AxiomInstCo con ind cos)
-  = do { unless (0 <= ind && ind < brListLength (coAxiomBranches con))
+  = do { unless (0 <= ind && ind < numBranches (coAxiomBranches con))
                 (bad_ax (ptext (sLit "index out of range")))
          -- See Note [Kind instantiation in coercions]
        ; let CoAxBranch { cab_tvs   = ktvs
diff --git a/compiler/iface/MkIface.hs b/compiler/iface/MkIface.hs
index 0bbd907464..a095ec84d2 100644
--- a/compiler/iface/MkIface.hs
+++ b/compiler/iface/MkIface.hs
@@ -1575,11 +1575,12 @@ coAxiomToIfaceDecl ax@(CoAxiom { co_ax_tc = tycon, co_ax_branches = branches
  = IfaceAxiom { ifName       = name
               , ifTyCon      = toIfaceTyCon tycon
               , ifRole       = role
-              , ifAxBranches = brListMap (coAxBranchToIfaceBranch tycon
-                                            (brListMap coAxBranchLHS branches))
-                                         branches }
+              , ifAxBranches = map (coAxBranchToIfaceBranch tycon
+                                     (map coAxBranchLHS branch_list))
+                                   branch_list }
  where
-   name = getOccName ax
+   branch_list = fromBranches branches
+   name        = getOccName ax
 
 -- 2nd parameter is the list of branch LHSs, for conversion from incompatible branches
 -- to incompatible indices
@@ -1679,7 +1680,7 @@ tyConToIfaceDecl env tycon
     to_if_fam_flav OpenSynFamilyTyCon        = IfaceOpenSynFamilyTyCon
     to_if_fam_flav (ClosedSynFamilyTyCon (Just ax))
       = IfaceClosedSynFamilyTyCon (Just (axn, ibr))
-      where defs = fromBranchList $ coAxiomBranches ax
+      where defs = fromBranches $ coAxiomBranches ax
             ibr  = map (coAxBranchToIfaceBranch' tycon) defs
             axn  = coAxiomName ax
     to_if_fam_flav (ClosedSynFamilyTyCon Nothing)
diff --git a/compiler/iface/TcIface.hs b/compiler/iface/TcIface.hs
index 5462fa29c2..5f91bad0e3 100644
--- a/compiler/iface/TcIface.hs
+++ b/compiler/iface/TcIface.hs
@@ -450,7 +450,7 @@ tc_iface_decl _ _ (IfaceAxiom { ifName = ax_occ, ifTyCon = tc
                              , co_ax_name     = tc_name
                              , co_ax_tc       = tc_tycon
                              , co_ax_role     = role
-                             , co_ax_branches = toBranchList tc_branches
+                             , co_ax_branches = manyBranches tc_branches
                              , co_ax_implicit = False }
        ; return (ACoAxiom axiom) }
 
diff --git a/compiler/typecheck/FamInst.hs b/compiler/typecheck/FamInst.hs
index 796bbcb582..7023a4c1f9 100644
--- a/compiler/typecheck/FamInst.hs
+++ b/compiler/typecheck/FamInst.hs
@@ -62,11 +62,7 @@ import Control.Arrow ( first, second )
 newFamInst :: FamFlavor -> CoAxiom Unbranched -> TcRnIf gbl lcl FamInst
 -- Freshen the type variables of the FamInst branches
 -- Called from the vectoriser monad too, hence the rather general type
-newFamInst flavor axiom@(CoAxiom { co_ax_branches = FirstBranch branch
-                                 , co_ax_tc = fam_tc })
-  | CoAxBranch { cab_tvs = tvs
-               , cab_lhs = lhs
-               , cab_rhs = rhs } <- branch
+newFamInst flavor axiom@(CoAxiom { co_ax_tc = fam_tc })
   = do { (subst, tvs') <- freshenTyVarBndrs tvs
        ; return (FamInst { fi_fam      = tyConName fam_tc
                          , fi_flavor   = flavor
@@ -75,6 +71,11 @@ newFamInst flavor axiom@(CoAxiom { co_ax_branches = FirstBranch branch
                          , fi_tys      = substTys subst lhs
                          , fi_rhs      = substTy  subst rhs
                          , fi_axiom    = axiom }) }
+  where
+    CoAxBranch { cab_tvs = tvs
+               , cab_lhs = lhs
+               , cab_rhs = rhs } = coAxiomSingleBranch axiom
+
 
 {-
 ************************************************************************
@@ -401,7 +402,7 @@ checkForInjectivityConflicts instEnvs famInst
     | isTypeFamilyTyCon tycon
     -- type family is injective in at least one argument
     , Injective inj <- familyTyConInjectivityInfo tycon = do
-    { let axiom = brFromUnbranchedSingleton (co_ax_branches (fi_axiom famInst))
+    { let axiom = coAxiomSingleBranch (fi_axiom famInst)
           conflicts = lookupFamInstEnvInjectivityConflicts inj instEnvs famInst
           -- see Note [Verifying injectivity annotation] in FamInstEnv
           errs = makeInjectivityErrors tycon axiom inj conflicts
diff --git a/compiler/typecheck/TcInteract.hs b/compiler/typecheck/TcInteract.hs
index 773f2ae6fc..49a5d4cc09 100644
--- a/compiler/typecheck/TcInteract.hs
+++ b/compiler/typecheck/TcInteract.hs
@@ -24,7 +24,7 @@ import PrelNames ( knownNatClassName, knownSymbolClassName,
                    callStackTyConKey, typeableClassName )
 import TysWiredIn ( ipClass, typeNatKind, typeSymbolKind )
 import Id( idType )
-import CoAxiom ( Eqn, CoAxiom(..), CoAxBranch(..), fromBranchList )
+import CoAxiom ( Eqn, CoAxiom(..), CoAxBranch(..), fromBranches )
 import Class
 import TyCon
 import DataCon( dataConWrapId )
@@ -1450,7 +1450,7 @@ improve_top_fun_eqs fam_envs fam_tc args rhs_ty
   | Just ax <- isClosedSynFamilyTyConWithAxiom_maybe fam_tc
   , Injective injective_args <- familyTyConInjectivityInfo fam_tc
   = concatMapM (injImproveEqns injective_args) $
-      buildImprovementData (fromBranchList (co_ax_branches ax))
+      buildImprovementData (fromBranches (co_ax_branches ax))
                            cab_lhs cab_rhs Just
 
   | otherwise
diff --git a/compiler/typecheck/TcRnDriver.hs b/compiler/typecheck/TcRnDriver.hs
index 8b4747575c..a7bfdd2655 100644
--- a/compiler/typecheck/TcRnDriver.hs
+++ b/compiler/typecheck/TcRnDriver.hs
@@ -1040,8 +1040,11 @@ checkBootTyCon tc1 tc2
     eqClosedFamilyAx (Just _) Nothing = False
     eqClosedFamilyAx (Just (CoAxiom { co_ax_branches = branches1 }))
                      (Just (CoAxiom { co_ax_branches = branches2 }))
-      =  brListLength branches1 == brListLength branches2
-      && (and $ brListZipWith eqClosedFamilyBranch branches1 branches2)
+      =  numBranches branches1 == numBranches branches2
+      && (and $ zipWith eqClosedFamilyBranch branch_list1 branch_list2)
+      where
+        branch_list1 = fromBranches branches1
+        branch_list2 = fromBranches branches2
 
     eqClosedFamilyBranch (CoAxBranch { cab_tvs = tvs1, cab_lhs = lhs1, cab_rhs = rhs1 })
                          (CoAxBranch { cab_tvs = tvs2, cab_lhs = lhs2, cab_rhs = rhs2 })
diff --git a/compiler/typecheck/TcSplice.hs b/compiler/typecheck/TcSplice.hs
index 52b52dbedf..2a21705c77 100644
--- a/compiler/typecheck/TcSplice.hs
+++ b/compiler/typecheck/TcSplice.hs
@@ -1125,7 +1125,8 @@ reifyTyCon tc
                              instances) }
          else do { eqns <-
                      case isClosedSynFamilyTyConWithAxiom_maybe tc of
-                       Just ax -> brListMapM reifyAxBranch $ coAxiomBranches ax
+                       Just ax -> mapM reifyAxBranch $
+                                  fromBranches $ coAxiomBranches ax
                        Nothing -> return []
                  ; return (TH.FamilyI
                       (TH.ClosedTypeFamilyD (reifyName tc) tvs' resultSig
diff --git a/compiler/typecheck/TcType.hs b/compiler/typecheck/TcType.hs
index 465efcca63..ffaef16cac 100644
--- a/compiler/typecheck/TcType.hs
+++ b/compiler/typecheck/TcType.hs
@@ -82,7 +82,7 @@ module TcType (
 
   ---------------------------------
   -- Predicate types
-  mkMinimalBySCs, transSuperClasses, transSuperClassesPred, 
+  mkMinimalBySCs, transSuperClasses, transSuperClassesPred,
   immSuperClasses,
   isImprovementPred,
 
@@ -1259,8 +1259,8 @@ occurCheckExpand dflags tv ty
     -- it and try again.
     go ty@(TyConApp tc tys)
       = case do { tys <- mapM go tys; return (mkTyConApp tc tys) } of
-          OC_OK ty 
-              | impredicative || isTauTyCon tc 
+          OC_OK ty
+              | impredicative || isTauTyCon tc
               -> return ty  -- First try to eliminate the tyvar from the args
               | otherwise
               -> OC_Forall  -- A type synonym with a forall on the RHS
@@ -1310,7 +1310,7 @@ Note [Kind polymorphic type classes]
     class C f where...   -- C :: forall k. k -> Constraint
     g :: forall (f::*). C f => f -> f
 
-Here the (C f) in the signature is really (C * f), and we 
+Here the (C f) in the signature is really (C * f), and we
 don't want to complain that the * isn't a type variable!
 -}
 
@@ -1331,7 +1331,7 @@ checkValidClsArgs flexible_contexts kts
   | otherwise         = all hasTyVarHead tys
   where
     (_, tys) = span isKind kts  -- see Note [Kind polymorphic type classes]
-   
+
 hasTyVarHead :: Type -> Bool
 -- Returns true of (a t1 .. tn), where 'a' is a type variable
 hasTyVarHead ty                 -- Haskell 98 allows predicates of form
@@ -1389,7 +1389,7 @@ immSuperClasses cls tys
 
 isImprovementPred :: PredType -> Bool
 -- Either it's an equality, or has some functional dependency
-isImprovementPred ty 
+isImprovementPred ty
   = case classifyPredType ty of
       EqPred NomEq t1 t2 -> not (t1 `tcEqType` t2)
       EqPred ReprEq _ _  -> False
@@ -1599,8 +1599,9 @@ orphNamesOfCoCon :: CoAxiom br -> NameSet
 orphNamesOfCoCon (CoAxiom { co_ax_tc = tc, co_ax_branches = branches })
   = orphNamesOfTyCon tc `unionNameSet` orphNamesOfCoAxBranches branches
 
-orphNamesOfCoAxBranches :: BranchList CoAxBranch br -> NameSet
-orphNamesOfCoAxBranches = brListFoldr (unionNameSet . orphNamesOfCoAxBranch) emptyNameSet
+orphNamesOfCoAxBranches :: Branches br -> NameSet
+orphNamesOfCoAxBranches
+  = foldr (unionNameSet . orphNamesOfCoAxBranch) emptyNameSet . fromBranches
 
 orphNamesOfCoAxBranch :: CoAxBranch -> NameSet
 orphNamesOfCoAxBranch (CoAxBranch { cab_lhs = lhs, cab_rhs = rhs })
@@ -1898,4 +1899,3 @@ size_type (ForAllTy _ ty)   = size_type ty
 
 sizeTypes :: [Type] -> TypeSize
 sizeTypes tys = sum (map sizeType tys)
-
diff --git a/compiler/typecheck/TcValidity.hs b/compiler/typecheck/TcValidity.hs
index ae416e76f9..b2a4f687b7 100644
--- a/compiler/typecheck/TcValidity.hs
+++ b/compiler/typecheck/TcValidity.hs
@@ -1224,9 +1224,10 @@ wrongATArgErr ty instTy =
 
 checkValidCoAxiom :: CoAxiom Branched -> TcM ()
 checkValidCoAxiom (CoAxiom { co_ax_tc = fam_tc, co_ax_branches = branches })
-  = do { _ <- brListMapM (checkValidCoAxBranch Nothing fam_tc) branches
-       ; brListFoldlM_ check_branch_compat [] branches }
+  = do { _ <- mapM (checkValidCoAxBranch Nothing fam_tc) branch_list
+       ; foldlM_ check_branch_compat [] branch_list }
   where
+    branch_list = fromBranches branches
     injectivity = familyTyConInjectivityInfo fam_tc
 
     check_branch_compat :: [CoAxBranch]    -- previous branches in reverse order
diff --git a/compiler/types/CoAxiom.hs b/compiler/types/CoAxiom.hs
index 66cec4c6ac..5b049a40f9 100644
--- a/compiler/types/CoAxiom.hs
+++ b/compiler/types/CoAxiom.hs
@@ -1,18 +1,16 @@
 -- (c) The University of Glasgow 2012
 
 {-# LANGUAGE CPP, DataKinds, DeriveDataTypeable, GADTs, KindSignatures,
-             ScopedTypeVariables, StandaloneDeriving #-}
+             ScopedTypeVariables, StandaloneDeriving, RoleAnnotations #-}
 
 -- | Module for coercion axioms, used to represent type family instances
 -- and newtypes
 
 module CoAxiom (
-       BranchFlag, Branched, Unbranched, BranchIndex, BranchList(..),
-       toBranchList, fromBranchList,
-       toBranchedList, toUnbranchedList,
-       brFromUnbranchedSingleton,
-       brListLength, brListNth, brListMap, brListFoldr, brListMapM,
-       brListFoldlM_, brListZipWith,
+       BranchFlag, Branched, Unbranched, BranchIndex, Branches,
+       manyBranches, unbranched,
+       fromBranches, numBranches,
+       mapAccumBranches,
 
        CoAxiom(..), CoAxBranch(..),
 
@@ -44,13 +42,15 @@ import BasicTypes
 import Data.Typeable ( Typeable )
 import SrcLoc
 import qualified Data.Data as Data
+import Data.Array
+import Data.List ( mapAccumL )
 
 #include "HsVersions.h"
 
 {-
 Note [Coercion axiom branches]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-In order to allow type family instance groups, an axiom needs to contain an
+In order to allow closed type families, an axiom needs to contain an
 ordered list of alternatives, called branches. The kind of the coercion built
 from an axiom is determined by which index is used when building the coercion
 from the axiom.
@@ -98,21 +98,21 @@ Note [Branched axioms]
 ~~~~~~~~~~~~~~~~~~~~~~
 Although a CoAxiom has the capacity to store many branches, in certain cases,
 we want only one. These cases are in data/newtype family instances, newtype
-coercions, and type family instances declared with "type instance ...", not
-"type instance where". Furthermore, these unbranched axioms are used in a
+coercions, and type family instances.
+Furthermore, these unbranched axioms are used in a
 variety of places throughout GHC, and it would difficult to generalize all of
 that code to deal with branched axioms, especially when the code can be sure
 of the fact that an axiom is indeed a singleton. At the same time, it seems
 dangerous to assume singlehood in various places through GHC.
 
 The solution to this is to label a CoAxiom with a phantom type variable
-declaring whether it is known to be a singleton or not. The list of branches
-is stored using a special form of list, declared below, that ensures that the
+declaring whether it is known to be a singleton or not. The branches
+are stored using a special datatype, declared below, that ensures that the
 type variable is accurate.
 
 ************************************************************************
 *                                                                      *
-                    Branch lists
+                    Branches
 *                                                                      *
 ************************************************************************
 -}
@@ -130,83 +130,47 @@ deriving instance Typeable 'Unbranched
 -- DataKinds and the promotion quote in client modules. This also means that
 -- we don't need to export the term-level constructors, which should never be used.
 
-data BranchList a (br :: BranchFlag) where
-  FirstBranch :: a -> BranchList a br
-  NextBranch :: a -> BranchList a br -> BranchList a Branched
-
-deriving instance Typeable BranchList
-
--- convert to/from lists
-toBranchList :: [a] -> BranchList a Branched
-toBranchList [] = pprPanic "toBranchList" empty
-toBranchList [b] = FirstBranch b
-toBranchList (h:t) = NextBranch h (toBranchList t)
-
-fromBranchList :: BranchList a br -> [a]
-fromBranchList (FirstBranch b) = [b]
-fromBranchList (NextBranch h t) = h : (fromBranchList t)
-
--- convert from any BranchList to a Branched BranchList
-toBranchedList :: BranchList a br -> BranchList a Branched
-toBranchedList (FirstBranch b) = FirstBranch b
-toBranchedList (NextBranch h t) = NextBranch h t
-
--- convert from any BranchList to an Unbranched BranchList
-toUnbranchedList :: BranchList a br -> BranchList a Unbranched
-toUnbranchedList (FirstBranch b) = FirstBranch b
-toUnbranchedList _ = pprPanic "toUnbranchedList" empty
-
--- Extract a singleton axiom from Unbranched BranchList
-brFromUnbranchedSingleton :: BranchList a Unbranched -> a
-brFromUnbranchedSingleton (FirstBranch b) = b
-
--- length
-brListLength :: BranchList a br -> Int
-brListLength (FirstBranch _) = 1
-brListLength (NextBranch _ t) = 1 + brListLength t
-
--- lookup
-brListNth :: BranchList a br -> BranchIndex -> a
-brListNth (FirstBranch b) 0 = b
-brListNth (NextBranch h _) 0 = h
-brListNth (NextBranch _ t) n = brListNth t (n-1)
-brListNth _ _ = pprPanic "brListNth" empty
-
--- map, fold
-brListMap :: (a -> b) -> BranchList a br -> [b]
-brListMap f (FirstBranch b) = [f b]
-brListMap f (NextBranch h t) = f h : (brListMap f t)
-
-brListFoldr :: (a -> b -> b) -> b -> BranchList a br -> b
-brListFoldr f x (FirstBranch b) = f b x
-brListFoldr f x (NextBranch h t) = f h (brListFoldr f x t)
-
-brListMapM :: Monad m => (a -> m b) -> BranchList a br -> m [b]
-brListMapM f (FirstBranch b) = f b >>= \fb -> return [fb]
-brListMapM f (NextBranch h t) = do { fh <- f h
-                                   ; ft <- brListMapM f t
-                                   ; return (fh : ft) }
-
-brListFoldlM_ :: forall a b m br. Monad m
-              => (a -> b -> m a) -> a -> BranchList b br -> m ()
-brListFoldlM_ f z brs = do { _ <- go z brs
-                           ; return () }
-  where go :: forall br'. a -> BranchList b br' -> m a
-        go acc (FirstBranch b)  = f acc b
-        go acc (NextBranch h t) = do { fh <- f acc h
-                                     ; go fh t }
-
--- zipWith
-brListZipWith :: (a -> b -> c) -> BranchList a br1 -> BranchList b br2 -> [c]
-brListZipWith f (FirstBranch a) (FirstBranch b) = [f a b]
-brListZipWith f (FirstBranch a) (NextBranch b _) = [f a b]
-brListZipWith f (NextBranch a _) (FirstBranch b) = [f a b]
-brListZipWith f (NextBranch a ta) (NextBranch b tb) = f a b : brListZipWith f ta tb
-
--- pretty-printing
-
-instance Outputable a => Outputable (BranchList a br) where
-  ppr = ppr . fromBranchList
+newtype Branches (br :: BranchFlag)
+  = MkBranches { unMkBranches :: Array BranchIndex CoAxBranch }
+  deriving Typeable
+type role Branches nominal
+
+manyBranches :: [CoAxBranch] -> Branches Branched
+manyBranches brs = ASSERT( snd bnds >= fst bnds )
+                   MkBranches (listArray bnds brs)
+  where
+    bnds = (0, length brs - 1)
+
+unbranched :: CoAxBranch -> Branches Unbranched
+unbranched br = MkBranches (listArray (0, 0) [br])
+
+toBranched :: Branches br -> Branches Branched
+toBranched = MkBranches . unMkBranches
+
+toUnbranched :: Branches br -> Branches Unbranched
+toUnbranched (MkBranches arr) = ASSERT( bounds arr == (0,0) )
+                                MkBranches arr
+
+fromBranches :: Branches br -> [CoAxBranch]
+fromBranches = elems . unMkBranches
+
+branchesNth :: Branches br -> BranchIndex -> CoAxBranch
+branchesNth (MkBranches arr) n = arr ! n
+
+numBranches :: Branches br -> Int
+numBranches (MkBranches arr) = snd (bounds arr) + 1
+
+-- | The @[CoAxBranch]@ passed into the mapping function is a list of
+-- all previous branches, reversed
+mapAccumBranches :: ([CoAxBranch] -> CoAxBranch -> CoAxBranch)
+                  -> Branches br -> Branches br
+mapAccumBranches f (MkBranches arr)
+  = MkBranches (listArray (bounds arr) (snd $ mapAccumL go [] (elems arr)))
+  where
+    go :: [CoAxBranch] -> CoAxBranch -> ([CoAxBranch], CoAxBranch)
+    go prev_branches cur_branch = ( cur_branch : prev_branches
+                                  , f prev_branches cur_branch )
+
 
 {-
 ************************************************************************
@@ -245,8 +209,7 @@ data CoAxiom br
     , co_ax_name     :: Name          -- name for pretty-printing
     , co_ax_role     :: Role          -- role of the axiom's equality
     , co_ax_tc       :: TyCon         -- the head of the LHS patterns
-    , co_ax_branches :: BranchList CoAxBranch br
-                                      -- the branches that form this axiom
+    , co_ax_branches :: Branches br   -- the branches that form this axiom
     , co_ax_implicit :: Bool          -- True <=> the axiom is "implicit"
                                       -- See Note [Implicit axioms]
          -- INVARIANT: co_ax_implicit == True implies length co_ax_branches == 1.
@@ -269,18 +232,18 @@ data CoAxBranch
 
 toBranchedAxiom :: CoAxiom br -> CoAxiom Branched
 toBranchedAxiom (CoAxiom unique name role tc branches implicit)
-  = CoAxiom unique name role tc (toBranchedList branches) implicit
+  = CoAxiom unique name role tc (toBranched branches) implicit
 
 toUnbranchedAxiom :: CoAxiom br -> CoAxiom Unbranched
 toUnbranchedAxiom (CoAxiom unique name role tc branches implicit)
-  = CoAxiom unique name role tc (toUnbranchedList branches) implicit
+  = CoAxiom unique name role tc (toUnbranched branches) implicit
 
 coAxiomNumPats :: CoAxiom br -> Int
 coAxiomNumPats = length . coAxBranchLHS . (flip coAxiomNthBranch 0)
 
 coAxiomNthBranch :: CoAxiom br -> BranchIndex -> CoAxBranch
 coAxiomNthBranch (CoAxiom { co_ax_branches = bs }) index
-  = brListNth bs index
+  = branchesNth bs index
 
 coAxiomArity :: CoAxiom br -> BranchIndex -> Arity
 coAxiomArity ax index
@@ -292,18 +255,19 @@ coAxiomName = co_ax_name
 coAxiomRole :: CoAxiom br -> Role
 coAxiomRole = co_ax_role
 
-coAxiomBranches :: CoAxiom br -> BranchList CoAxBranch br
+coAxiomBranches :: CoAxiom br -> Branches br
 coAxiomBranches = co_ax_branches
 
 coAxiomSingleBranch_maybe :: CoAxiom br -> Maybe CoAxBranch
-coAxiomSingleBranch_maybe (CoAxiom { co_ax_branches = branches })
-  | FirstBranch br <- branches
-  = Just br
+coAxiomSingleBranch_maybe (CoAxiom { co_ax_branches = MkBranches arr })
+  | snd (bounds arr) == 0
+  = Just $ arr ! 0
   | otherwise
   = Nothing
 
 coAxiomSingleBranch :: CoAxiom Unbranched -> CoAxBranch
-coAxiomSingleBranch (CoAxiom { co_ax_branches = FirstBranch br }) = br
+coAxiomSingleBranch (CoAxiom { co_ax_branches = MkBranches arr })
+  = arr ! 0
 
 coAxiomTyCon :: CoAxiom br -> TyCon
 coAxiomTyCon = co_ax_tc
diff --git a/compiler/types/Coercion.hs b/compiler/types/Coercion.hs
index 5d4329d5ff..b73ca4969b 100644
--- a/compiler/types/Coercion.hs
+++ b/compiler/types/Coercion.hs
@@ -756,7 +756,7 @@ ppr_forall_co p ty
 pprCoAxiom :: CoAxiom br -> SDoc
 pprCoAxiom ax@(CoAxiom { co_ax_tc = tc, co_ax_branches = branches })
   = hang (ptext (sLit "axiom") <+> ppr ax <+> dcolon)
-       2 (vcat (map (pprCoAxBranch tc) $ fromBranchList branches))
+       2 (vcat (map (pprCoAxBranch tc) $ fromBranches branches))
 
 pprCoAxBranch :: TyCon -> CoAxBranch -> SDoc
 pprCoAxBranch fam_tc (CoAxBranch { cab_tvs = tvs
@@ -1215,7 +1215,7 @@ mkNewTypeCo name tycon tvs roles rhs_ty
             , co_ax_implicit = True  -- See Note [Implicit axioms] in TyCon
             , co_ax_role     = Representational
             , co_ax_tc       = tycon
-            , co_ax_branches = FirstBranch branch }
+            , co_ax_branches = unbranched branch }
   where branch = CoAxBranch { cab_loc     = getSrcSpan name
                             , cab_tvs     = tvs
                             , cab_lhs     = mkTyVarTys tvs
diff --git a/compiler/types/FamInstEnv.hs b/compiler/types/FamInstEnv.hs
index 63d76c4084..a41e4539c7 100644
--- a/compiler/types/FamInstEnv.hs
+++ b/compiler/types/FamInstEnv.hs
@@ -60,6 +60,7 @@ import Pair
 import SrcLoc
 import NameSet
 import FastString
+import Data.Function ( on )
 
 {-
 ************************************************************************
@@ -250,10 +251,9 @@ mkImportedFamInst fam mb_tcs axiom
       fi_flavor = flavor }
   where
      -- See Note [Lazy axiom match]
-     ~(CoAxiom { co_ax_branches =
-       ~(FirstBranch ~(CoAxBranch { cab_lhs = tys
-                                  , cab_tvs = tvs
-                                  , cab_rhs = rhs })) }) = axiom
+     ~(CoAxBranch { cab_lhs = tys
+                  , cab_tvs = tvs
+                  , cab_rhs = rhs }) = coAxiomSingleBranch axiom
 
          -- Derive the flavor for an imported FamInst rather disgustingly
          -- Maybe we should store it in the IfaceFamInst?
@@ -353,7 +353,7 @@ familyInstances (pkg_fie, home_fie) fam
 -- Used in the implementation of ":info" in GHCi.
 orphNamesOfFamInst :: FamInst -> NameSet
 orphNamesOfFamInst fam_inst
-  = orphNamesOfTypes (concat (brListMap cab_lhs (coAxiomBranches axiom)))
+  = orphNamesOfTypes (concat (map cab_lhs (fromBranches $ coAxiomBranches axiom)))
     `extendNameSet` getName (coAxiomTyCon axiom)
   where
     axiom = fi_axiom fam_inst
@@ -382,8 +382,8 @@ identicalFamInstHead :: FamInst -> FamInst -> Bool
 -- Used for overriding in GHCi
 identicalFamInstHead (FamInst { fi_axiom = ax1 }) (FamInst { fi_axiom = ax2 })
   =  coAxiomTyCon ax1 == coAxiomTyCon ax2
-  && brListLength brs1 == brListLength brs2
-  && and (brListZipWith identical_branch brs1 brs2)
+  && numBranches brs1 == numBranches brs2
+  && and ((zipWith identical_branch `on` fromBranches) brs1 brs2)
   where
     brs1 = coAxiomBranches ax1
     brs2 = coAxiomBranches ax2
@@ -528,14 +528,10 @@ injectiveBranches injectivity
 -- See Note [Storing compatibility] in CoAxiom
 computeAxiomIncomps :: CoAxiom br -> CoAxiom br
 computeAxiomIncomps ax@(CoAxiom { co_ax_branches = branches })
-  = ax { co_ax_branches = go [] branches }
+  = ax { co_ax_branches = mapAccumBranches go branches }
   where
-    go :: [CoAxBranch] -> BranchList CoAxBranch br -> BranchList CoAxBranch br
-    go prev_branches (FirstBranch br)
-      = FirstBranch (br { cab_incomps = mk_incomps br prev_branches })
-    go prev_branches (NextBranch br tail)
-      = let br' = br { cab_incomps = mk_incomps br prev_branches } in
-        NextBranch br' (go (br' : prev_branches) tail)
+    go :: [CoAxBranch] -> CoAxBranch -> CoAxBranch
+    go prev_branches br = br { cab_incomps = mk_incomps br prev_branches }
 
     mk_incomps :: CoAxBranch -> [CoAxBranch] -> [CoAxBranch]
     mk_incomps br = filter (not . compatibleBranches br)
@@ -583,7 +579,7 @@ mkBranchedCoAxiom ax_name fam_tc branches
             , co_ax_tc       = fam_tc
             , co_ax_role     = Nominal
             , co_ax_implicit = False
-            , co_ax_branches = toBranchList branches }
+            , co_ax_branches = manyBranches branches }
 
 mkUnbranchedCoAxiom :: Name -> TyCon -> CoAxBranch -> CoAxiom Unbranched
 mkUnbranchedCoAxiom ax_name fam_tc branch
@@ -592,7 +588,7 @@ mkUnbranchedCoAxiom ax_name fam_tc branch
             , co_ax_tc       = fam_tc
             , co_ax_role     = Nominal
             , co_ax_implicit = False
-            , co_ax_branches = FirstBranch (branch { cab_incomps = [] }) }
+            , co_ax_branches = unbranched (branch { cab_incomps = [] }) }
 
 mkSingleCoAxiom :: Role -> Name
                 -> [TyVar] -> TyCon -> [Type] -> Type
@@ -606,7 +602,7 @@ mkSingleCoAxiom role ax_name tvs fam_tc lhs_tys rhs_ty
             , co_ax_tc       = fam_tc
             , co_ax_role     = role
             , co_ax_implicit = False
-            , co_ax_branches = FirstBranch (branch { cab_incomps = [] }) }
+            , co_ax_branches = unbranched (branch { cab_incomps = [] }) }
   where
     branch = mkCoAxBranch tvs lhs_tys rhs_ty (getSrcSpan ax_name)
 
@@ -815,7 +811,7 @@ lookupFamInstEnvInjectivityConflicts injList (pkg_ie, home_ie)
 
       lookup_inj_fam_conflicts ie
           | isOpenFamilyTyCon fam, Just (FamIE insts) <- lookupUFM ie fam
-          = map (brFromUnbranchedSingleton . co_ax_branches . fi_axiom) $
+          = map (coAxiomSingleBranch . fi_axiom) $
             filter isInjConflict insts
           | otherwise = []
 
@@ -1017,7 +1013,7 @@ chooseBranch axiom tys
   = do { let num_pats = coAxiomNumPats axiom
              (target_tys, extra_tys) = splitAt num_pats tys
              branches = coAxiomBranches axiom
-       ; (ind, inst_tys) <- findBranch (fromBranchList branches) target_tys
+       ; (ind, inst_tys) <- findBranch (fromBranches branches) target_tys
        ; return (ind, inst_tys ++ extra_tys) }
 
 -- The axiom must *not* be oversaturated