From f08b39b6c5f8f5a8d181373cea45f858e2bd6473 Mon Sep 17 00:00:00 2001
From: Yuchen Pei <hi@ypei.me>
Date: Tue, 16 Aug 2022 09:56:48 +1000
Subject: handling associated types for classes

Also fixed the familydecl a bit
---
 haddock-api/src/Haddock/Backends/Org.hs | 72 +++++++++++++++++++++------------
 1 file changed, 47 insertions(+), 25 deletions(-)

diff --git a/haddock-api/src/Haddock/Backends/Org.hs b/haddock-api/src/Haddock/Backends/Org.hs
index 9d02d0db..d83d6319 100644
--- a/haddock-api/src/Haddock/Backends/Org.hs
+++ b/haddock-api/src/Haddock/Backends/Org.hs
@@ -534,46 +534,68 @@ ppTyClDecl (ClassDecl {..}) docs subdocs path level =
     ++ ppDocForDecl docs (Just level)
     -- TODO: do we need an aDoc here instead of M.empty?
     -- TODO: handle default sigs
+    ++ concatMap
+         ((\assoc -> ppFamilyDecl assoc False emptyDoc subdocs path (level + 1))
+         . unLoc
+         )
+         tcdATs
     ++ concatMap
          ((\sig -> ppSig sig emptyDoc subdocs path (level + 1)) . unLoc)
          tcdSigs
 -- type family ... where
 -- TODO: handle infix
-ppTyClDecl (FamDecl _ (FamilyDecl _ (ClosedTypeFamily mbEqns) TopLevel (L _ name) tyvars _fixity (L _ resSig) mbInj)) docs subdocs path level
+ppTyClDecl (FamDecl _ familyDecl) docs subdocs path level =
+  ppFamilyDecl familyDecl True docs subdocs path level
+
+ppFamilyDecl
+  :: FamilyDecl DocNameI
+  -> Bool
+  -> DocForDecl DocName
+  -> SubDocs
+  -> ModPath
+  -> Int
+  -> [OrgBlock]
+ppFamilyDecl (FamilyDecl _ info@(ClosedTypeFamily mbEqns) TopLevel (L _ name) tyvars _fixity (L _ resSig) mbInj) isFamDecl docs subdocs path level
   = Heading
       level
-      (  [plaintext "type family ", Plain $ docNameToDoc name, Whitespace]
+      (  ppFamilyInfo info isFamDecl
+      ++ [Whitespace, Plain $ docNameToDoc name, Whitespace]
       ++ ppLHsQTyVars tyvars
-      ++ ppFamilyResultSig resSig "="
+      ++ ppFamilyResultSig resSig
       ++ maybe [] ppLInjectivityAnn mbInj
       ++ [plaintext " where"]
       )
       (cIdPaths path name)
-    :  ppDocForDecl docs (Just level)
+    :  (if isFamDecl
+         then ppDocForDecl docs (Just level)
+         else maybe [] (`ppDocForDecl` (Just level)) (lookup name subdocs)
+       )
     ++ concatMap (\x -> ppLTyFamInstEqn x subdocs path (level + 1))
                  (fromMaybe [] mbEqns)
--- data family
--- type family
--- DataFamily or OpenTypeFamily
-ppTyClDecl (FamDecl _ (FamilyDecl _ info TopLevel (L _ name) tyvars _fixity (L _ resSig) mbInj)) docs _ path level
+ppFamilyDecl (FamilyDecl _ info TopLevel (L _ name) tyvars _fixity (L _ resSig) mbInj) isFamDecl docs subdocs path level
   = Heading
       level
-      (  [pre, Plain $ docNameToDoc name, Whitespace]
+      (  ppFamilyInfo info isFamDecl
+      ++ [Whitespace, Plain $ docNameToDoc name, Whitespace]
       ++ ppLHsQTyVars tyvars
-      ++ ppFamilyResultSig resSig op
+      ++ ppFamilyResultSig resSig
       ++ maybe [] ppLInjectivityAnn mbInj
       )
       (cIdPaths path name)
-    : ppDocForDecl docs (Just level)
+    : (if isFamDecl
+        then ppDocForDecl docs (Just level)
+        else maybe [] (`ppDocForDecl` (Just level)) (lookup name subdocs)
+      )
+ppFamilyDecl _ _ docs _ _ level =
+  unimpHeading "FamilyDecl" level : ppDocForDecl docs (Just level)
+
+ppFamilyInfo :: FamilyInfo DocNameI -> Bool -> [OrgInline]
+ppFamilyInfo info isFamDecl = dataOrType : family
  where
-  pre = case info of
-    DataFamily     -> plaintext "data family "
-    OpenTypeFamily -> plaintext "type family "
-  op = case info of
-    DataFamily -> "::"
-    _          -> "="
-ppTyClDecl (FamDecl{}) docs _ _ level =
-  unimpHeading "FamDecl" level : ppDocForDecl docs (Just level)
+  dataOrType = case info of
+    DataFamily -> plaintext "data"
+    _          -> plaintext "type"
+  family = if isFamDecl then [plaintext " family"] else []
 
 ppLTyFamInstEqn
   :: LTyFamInstEqn DocNameI -> SubDocs -> ModPath -> Int -> [OrgBlock]
@@ -604,12 +626,12 @@ ppLInjectivityAnn (L _ (InjectivityAnn _ (L _ l) rs)) =
   ]
 ppLInjectivityAnn _ = [plaintext $ unimp "LInjectivityAnn"]
 
-ppFamilyResultSig :: FamilyResultSig DocNameI -> String -> [OrgInline]
-ppFamilyResultSig (KindSig _ (L _ x)) op =
-  [Whitespace, plaintext op, Whitespace] ++ ppHsType x
-ppFamilyResultSig (NoSig{}) _ = []
-ppFamilyResultSig (TyVarSig _ x) op =
-  [Whitespace, plaintext op, Whitespace] ++ ppLHsTyVarBndr x
+ppFamilyResultSig :: FamilyResultSig DocNameI -> [OrgInline]
+ppFamilyResultSig (KindSig _ (L _ x)) =
+  [Whitespace, plaintext "::", Whitespace] ++ ppHsType x
+ppFamilyResultSig (NoSig{}) = []
+ppFamilyResultSig (TyVarSig _ x) =
+  [Whitespace, plaintext "=", Whitespace] ++ ppLHsTyVarBndr x
 
 ppDataDefn :: HsDataDefn DocNameI -> SubDocs -> ModPath -> Int -> [OrgBlock]
 ppDataDefn (HsDataDefn _ _ _ _ _ cons _derivs) subdocs path level =
-- 
cgit v1.2.3