diff --git a/TESTS.md b/TESTS.md
--- a/TESTS.md
+++ b/TESTS.md
@@ -1786,3 +1786,23 @@
   (\(_ :: ErrorCall) -> pure 2)
 
 ```
+
+lippirk Comments on functions in where clause not quite right #540
+
+```haskell
+-- https://github.com/chrisdone/hindent/issues/540
+topLevelFunc1 = f
+  where
+    -- comment on func in where clause
+    -- stays in the where clause
+    f = undefined
+
+topLevelFunc2 = f . g
+  where
+    {- multi
+       line
+       comment -}
+    f = undefined
+    -- single line comment
+    g = undefined
+```
diff --git a/hindent.cabal b/hindent.cabal
--- a/hindent.cabal
+++ b/hindent.cabal
@@ -1,5 +1,5 @@
 name:                hindent
-version:             5.3.0
+version:             5.3.1
 synopsis:            Extensible Haskell pretty printer
 description:         Extensible Haskell pretty printer. Both a library and an executable.
                      .
diff --git a/src/HIndent.hs b/src/HIndent.hs
--- a/src/HIndent.hs
+++ b/src/HIndent.hs
@@ -1,6 +1,6 @@
 {-# OPTIONS_GHC -fno-warn-name-shadowing #-}
 {-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE OverloadedStrings, TupleSections, ScopedTypeVariables, PatternGuards #-}
+{-# LANGUAGE OverloadedStrings, ScopedTypeVariables, PatternGuards #-}
 
 -- | Haskell indenter.
 
@@ -32,6 +32,7 @@
 import qualified Data.ByteString.UTF8 as UTF8
 import qualified Data.ByteString.Unsafe as S
 import           Data.Char
+import           Data.Foldable (foldr')
 import           Data.Either
 import           Data.Function
 import           Data.Functor.Identity
@@ -178,7 +179,7 @@
               , psLine = 1
               , psConfig = config
               , psInsideCase = False
-              , psHardLimit = False
+              , psFitOnOneLine = False
               , psEolComment = False
               }))))
 
@@ -299,17 +300,16 @@
     (traverseBackwards
      -- Finally, collect backwards comments which come after each node.
        (collectCommentsBy
-          (<>)
           CommentAfterLine
           (\nodeSpan commentSpan ->
               fst (srcSpanStart commentSpan) >= fst (srcSpanEnd nodeSpan)))) <=<
+  shortCircuit addCommentsToTopLevelWhereClauses <=<
   shortCircuit
     (traverse
      -- Collect forwards comments which start at the end line of a
      -- node: Does the start line of the comment match the end-line
      -- of the node?
        (collectCommentsBy
-          (<>)
           CommentSameLine
           (\nodeSpan commentSpan ->
               fst (srcSpanStart commentSpan) == fst (srcSpanEnd nodeSpan)))) <=<
@@ -319,7 +319,6 @@
      -- node: Does the start line & end line of the comment match
      -- that of the node?
        (collectCommentsBy
-          (<>)
           CommentSameLine
           (\nodeSpan commentSpan ->
               fst (srcSpanStart commentSpan) == fst (srcSpanStart nodeSpan) &&
@@ -329,7 +328,6 @@
      -- First, collect forwards comments for declarations which both
      -- start on column 1 and occur before the declaration.
        (collectCommentsBy
-          (<>)
           CommentBeforeLine
           (\nodeSpan commentSpan ->
               (snd (srcSpanStart nodeSpan) == 1 &&
@@ -353,34 +351,84 @@
 -- comment means to remove it from the pool of available comments in
 -- the State. This allows for a multiple pass approach.
 collectCommentsBy
-  :: ([NodeComment] -> [NodeComment] -> [NodeComment])
-  -> (SrcSpan -> SomeComment -> NodeComment)
+  :: (SrcSpan -> SomeComment -> NodeComment)
   -> (SrcSpan -> SrcSpan -> Bool)
   -> NodeInfo
   -> State [Comment] NodeInfo
-collectCommentsBy append cons predicate nodeInfo@(NodeInfo (SrcSpanInfo nodeSpan _) _) = do
+collectCommentsBy cons predicate nodeInfo@(NodeInfo (SrcSpanInfo nodeSpan _) _) = do
   comments <- get
   let (others, mine) =
         partitionEithers
           (map
-             (\comment@(Comment multiLine commentSpan commentString) ->
-                 if predicate nodeSpan (setFilename commentString commentSpan)
-                   then Right
-                          (cons
-                             commentSpan
-                             ((if multiLine
-                                 then MultiLine
-                                 else EndOfLine)
-                                commentString))
+             (\comment@(Comment _ commentSpan _) ->
+                 if predicate nodeSpan commentSpan
+                   then Right comment
                    else Left comment)
              comments)
   put others
-  return
-    (nodeInfo
-     { nodeInfoComments = append (nodeInfoComments nodeInfo) mine
-     })
+  return $ addCommentsToNode cons mine nodeInfo
+
+-- | Reintroduce comments which were immediately above declarations in where clauses.
+-- Affects where clauses of top level declarations only.
+addCommentsToTopLevelWhereClauses ::
+     Module NodeInfo -> State [Comment] (Module NodeInfo)
+addCommentsToTopLevelWhereClauses (Module x x' x'' x''' topLevelDecls) =
+  Module x x' x'' x''' <$>
+  traverse addCommentsToWhereClauses topLevelDecls
   where
-    setFilename cs sp =
-      sp
-      { srcSpanFilename = cs
-      }
+    addCommentsToWhereClauses ::
+         Decl NodeInfo -> State [Comment] (Decl NodeInfo)
+    addCommentsToWhereClauses (PatBind x x' x'' (Just (BDecls x''' whereDecls))) = do
+      newWhereDecls <- traverse addCommentsToPatBind whereDecls
+      return $ PatBind x x' x'' (Just (BDecls x''' newWhereDecls))
+    addCommentsToWhereClauses other = return other
+    addCommentsToPatBind :: Decl NodeInfo -> State [Comment] (Decl NodeInfo)
+    addCommentsToPatBind (PatBind bindInfo (PVar x (Ident declNodeInfo declString)) x' x'') = do
+      bindInfoWithComments <- addCommentsBeforeNode bindInfo
+      return $
+        PatBind
+          bindInfoWithComments
+          (PVar x (Ident declNodeInfo declString))
+          x'
+          x''
+    addCommentsToPatBind other = return other
+    addCommentsBeforeNode :: NodeInfo -> State [Comment] NodeInfo
+    addCommentsBeforeNode nodeInfo = do
+      comments <- get
+      let (notAbove, above) = partitionAboveNotAbove comments nodeInfo
+      put notAbove
+      return $ addCommentsToNode CommentBeforeLine above nodeInfo
+    partitionAboveNotAbove :: [Comment] -> NodeInfo -> ([Comment], [Comment])
+    partitionAboveNotAbove cs (NodeInfo (SrcSpanInfo nodeSpan _) _) =
+      fst $
+      foldr'
+        (\comment@(Comment _ commentSpan _) ((ls, rs), lastSpan) ->
+           if comment `isAbove` lastSpan
+             then ((ls, comment : rs), commentSpan)
+             else ((comment : ls, rs), lastSpan))
+        (([], []), nodeSpan)
+        cs
+    isAbove :: Comment -> SrcSpan -> Bool
+    isAbove (Comment _ commentSpan _) span =
+      let (_, commentColStart) = srcSpanStart commentSpan
+          (commentLnEnd, _) = srcSpanEnd commentSpan
+          (lnStart, colStart) = srcSpanStart span
+       in commentColStart == colStart && commentLnEnd + 1 == lnStart
+addCommentsToTopLevelWhereClauses other = return other
+
+addCommentsToNode :: (SrcSpan -> SomeComment -> NodeComment)
+                  -> [Comment]
+                  -> NodeInfo
+                  -> NodeInfo
+addCommentsToNode mkNodeComment newComments nodeInfo@(NodeInfo (SrcSpanInfo _ _) existingComments) =
+  nodeInfo
+    {nodeInfoComments = existingComments <> map mkBeforeNodeComment newComments}
+  where
+    mkBeforeNodeComment :: Comment -> NodeComment
+    mkBeforeNodeComment (Comment multiLine commentSpan commentString) =
+      mkNodeComment
+        commentSpan
+        ((if multiLine
+            then MultiLine
+            else EndOfLine)
+           commentString)
diff --git a/src/HIndent/Pretty.hs b/src/HIndent/Pretty.hs
--- a/src/HIndent/Pretty.hs
+++ b/src/HIndent/Pretty.hs
@@ -231,7 +231,7 @@
 write :: String -> Printer ()
 write x =
   do eol <- gets psEolComment
-     hardFail <- gets psHardLimit
+     hardFail <- gets psFitOnOneLine
      let addingNewline = eol && x /= "\n"
      when addingNewline newline
      state <- get
@@ -1143,7 +1143,11 @@
                        (pretty d))
 
 instance Pretty GadtDecl where
+#if MIN_VERSION_haskell_src_exts(1,21,0)
+  prettyInternal (GadtDecl _ name _ _ fields t) =
+#else
   prettyInternal (GadtDecl _ name fields t) =
+#endif
     horVar `ifFitsOnOneLineOrElse` verVar
     where
       fields' p =
@@ -1442,8 +1446,10 @@
 instance Pretty FunDep where
   prettyInternal = pretty'
 
+#if !MIN_VERSION_haskell_src_exts(1,21,0)
 instance Pretty Kind where
   prettyInternal = pretty'
+#endif
 
 instance Pretty ResultSig where
   prettyInternal (KindSig _ kind) = pretty kind
@@ -1842,6 +1848,9 @@
          pretty n
 typ (TyQuasiQuote _ n s) = quotation n (string s)
 typ (TyUnboxedSum{}) = error "FIXME: No implementation for TyUnboxedSum."
+#if MIN_VERSION_haskell_src_exts(1,21,0)
+typ (TyStar _) = write "*"
+#endif
 
 prettyTopName :: Name NodeInfo -> Printer ()
 prettyTopName x@Ident{} = pretty x
@@ -2030,26 +2039,28 @@
 fitsOnOneLine :: Printer a -> Printer (Maybe PrintState)
 fitsOnOneLine p =
   do st <- get
-     put st { psHardLimit = True}
+     put st { psFitOnOneLine = True}
      ok <- fmap (const True) p <|> return False
      st' <- get
      put st
+     guard $ ok || not (psFitOnOneLine st)
      return (if ok
-                then Just st' { psHardLimit = psHardLimit st }
+                then Just st' { psFitOnOneLine = psFitOnOneLine st }
                 else Nothing)
 
 -- | If first printer fits, use it, else use the second one.
 ifFitsOnOneLineOrElse :: Printer a -> Printer a -> Printer a
 ifFitsOnOneLineOrElse a b = do
   stOrig <- get
-  put stOrig{psHardLimit = True}
+  put stOrig{psFitOnOneLine = True}
   res <- fmap Just a <|> return Nothing
   case res of
     Just r -> do
-      modify $ \st -> st{psHardLimit = psHardLimit stOrig}
+      modify $ \st -> st{psFitOnOneLine = psFitOnOneLine stOrig}
       return r
     Nothing -> do
       put stOrig
+      guard $ not (psFitOnOneLine stOrig)
       b
 
 bindingGroup :: Binds NodeInfo -> Printer ()
diff --git a/src/HIndent/Types.hs b/src/HIndent/Types.hs
--- a/src/HIndent/Types.hs
+++ b/src/HIndent/Types.hs
@@ -52,8 +52,9 @@
     -- ^ Configuration of max colums and indentation style.
   , psInsideCase :: !Bool
     -- ^ Whether we're in a case statement, used for Rhs printing.
-  , psHardLimit :: !Bool
-    -- ^ Bail out if we exceed current column.
+  , psFitOnOneLine :: !Bool
+    -- ^ Bail out if we need to print beyond the current line or
+    -- the maximum column.
   , psEolComment :: !Bool
   }