[New Multicolumn] Add support for column-span:all

Source/core/rendering/RenderMultiColumnFlowThread.cpp

 /*
  * Copyright (C) 2012 Apple Inc.  All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions
  * are met:
  * 1. Redistributions of source code must retain the above copyright
  *    notice, this list of conditions and the following disclaimer.
  * 2. Redistributions in binary form must reproduce the above copyright
  *    notice, this list of conditions and the following disclaimer in the
  *    documentation and/or other materials provided with the distribution.
  *
  * THIS SOFTWARE IS PROVIDED BY APPLE INC. ``AS IS'' AND ANY
  * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL APPLE COMPUTER, INC. OR
  * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
  * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
  * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
  * PROFITS; OR BUSINESS IN..0TERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
  * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  * OF THIS SOFTWARE, EVEN IF  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
 #include "config.h"
 #include "core/rendering/RenderMultiColumnFlowThread.h"
 
 #include "core/rendering/RenderMultiColumnSet.h"
+#include "core/rendering/RenderMultiColumnSpannerSet.h"
 
 namespace blink {
 
 RenderMultiColumnFlowThread::RenderMultiColumnFlowThread()
-    : m_columnCount(1)
+    : m_lastSetWorkedOn(0)
+    , m_columnCount(1)
     , m_columnHeightAvailable(0)
     , m_inBalancingPass(false)
     , m_needsColumnHeightsRecalculation(false)
     , m_progressionIsInline(true)
+    , m_beingEvacuated(false)
 {
     setFlowThreadState(InsideInFlowThread);
 }
 
 RenderMultiColumnFlowThread::~RenderMultiColumnFlowThread()
 {
 }
 
 RenderMultiColumnFlowThread* RenderMultiColumnFlowThread::createAnonymous(Document& document, RenderStyle* parentStyle)
 {
@@ skipping 14 matching lines @@
 
 RenderMultiColumnSet* RenderMultiColumnFlowThread::lastMultiColumnSet() const
 {
     for (RenderObject* sibling = multiColumnBlockFlow()->lastChild(); sibling; sibling = sibling->previousSibling()) {
         if (sibling->isRenderMultiColumnSet())
             return toRenderMultiColumnSet(sibling);
     }
     return 0;
 }
 
-void RenderMultiColumnFlowThread::addChild(RenderObject* newChild, RenderObject* beforeChild)
+RenderMultiColumnSet* RenderMultiColumnFlowThread::findSetRendering(RenderObject* renderer) const
 {
-    RenderBlockFlow::addChild(newChild, beforeChild);
-    if (firstMultiColumnSet())
-        return;
+    ASSERT(!containingColumnSpanner(renderer)); // should not be used for spanners or content inside them.
+    RenderMultiColumnSet* multicolSet = firstMultiColumnSet();
+    if (!multicolSet)
+        return 0;
 
-    // For now we only create one column set. It's created as soon as the multicol container gets
-    // any content at all.
-    RenderMultiColumnSet* newSet = RenderMultiColumnSet::createAnonymous(this, multiColumnBlockFlow()->style());
+    if (!multicolSet->nextSiblingMultiColumnSet()) {
+        // There is only one set. This is easy, then: if it's in the flow thread, it's part of the set.
+        return renderer->isDescendantOf(this) ? multicolSet : 0;
+    }
 
-    // Need to skip RenderBlockFlow's implementation of addChild(), or we'd get redirected right
-    // back here.
-    multiColumnBlockFlow()->RenderBlock::addChild(newSet);
+    // This is SLOW! But luckily very uncommon. You need to dynamically insert a spanner into the
+    // middle of the tree to trigger this.
+    for (; multicolSet; multicolSet = multicolSet->nextSiblingMultiColumnSet()) {
+        if (multicolSet->isRenderMultiColumnSpannerSet())
+            continue;
+        RenderObject* startRenderer;
+        if (RenderMultiColumnSet* sibling = multicolSet->previousSiblingMultiColumnSet()) {
+            // Adjacent sets should not occur. Currently we would have no way of figuring out what each
+            // of them contains then.
+            ASSERT(sibling->isRenderMultiColumnSpannerSet());
+            startRenderer = toRenderMultiColumnSpannerSet(sibling)->rendererInFlowThread()->nextInPreOrderAfterChildren(this);
+        } else {
+            startRenderer = firstChild();
+        }
+        ASSERT(startRenderer);
 
-    invalidateRegions();
+        RenderObject* stopRenderer;
+        if (RenderMultiColumnSet* sibling = multicolSet->nextSiblingMultiColumnSet()) {
+            // Adjacent sets should not occur. Currently we would have no way of figuring out what each
+            // of them contains then.
+            ASSERT(sibling->isRenderMultiColumnSpannerSet());
+            stopRenderer = toRenderMultiColumnSpannerSet(sibling)->rendererInFlowThread();
+        } else {
+            stopRenderer = 0;
+        }
+
+        for (RenderObject* walker = startRenderer; walker != stopRenderer; walker = walker->nextInPreOrder(this)) {
+            if (walker == renderer)
+                return multicolSet;
+        }
+    }
+
+    return 0;
+}
+
+RenderMultiColumnSpannerSet* RenderMultiColumnFlowThread::containingColumnSpanner(const RenderObject* descendant) const
+{
+    ASSERT(descendant->isDescendantOf(this));
+
+    // Before we spend time on searching the ancestry, see if there's a quick way to determine
+    // whether there might be any spanners at all.
+    RenderMultiColumnSet* firstSet = firstMultiColumnSet();
+    if (!firstSet || (firstSet == lastMultiColumnSet() && !firstSet->isRenderMultiColumnSpannerSet()))
+        return 0;
+
+    // We have spanners. See if the renderer in question is one or inside of one then.
+    for (const RenderObject* ancestor = descendant; ancestor && ancestor != this; ancestor = ancestor->parent()) {
+        if (RenderMultiColumnSpannerSet* spanner = m_spannerMap.get(ancestor))
+            return spanner;
+    }
+    return 0;
 }
 
 void RenderMultiColumnFlowThread::populate()
 {
     RenderBlockFlow* multicolContainer = multiColumnBlockFlow();
     ASSERT(!nextSibling());
     // Reparent children preceding the flow thread into the flow thread. It's multicol content
     // now. At this point there's obviously nothing after the flow thread, but renderers (column
     // sets and spanners) will be inserted there as we insert elements into the flow thread.
     multicolContainer->moveChildrenTo(this, multicolContainer->firstChild(), this, true);
 }
 
 void RenderMultiColumnFlowThread::evacuateAndDestroy()
 {
     RenderBlockFlow* multicolContainer = multiColumnBlockFlow();
+    m_beingEvacuated = true;
 
     // Remove all sets.
     while (RenderMultiColumnSet* columnSet = firstMultiColumnSet())
         columnSet->destroy();
 
     ASSERT(!previousSibling());
     ASSERT(!nextSibling());
 
     // Finally we can promote all flow thread's children. Before we move them to the flow thread's
     // container, we need to unregister the flow thread, so that they aren't just re-added again to
@@ skipping 47 matching lines @@
         m_needsColumnHeightsRecalculation = false;
         return;
     }
 
     for (RenderMultiColumnSet* columnSet = firstMultiColumnSet(); columnSet; columnSet = columnSet->nextSiblingMultiColumnSet()) {
         if (!m_inBalancingPass) {
             // This is the initial layout pass. We need to reset the column height, because contents
             // typically have changed.
             columnSet->resetColumnHeight();
         }
+        columnSet->resetFlow();
     }
 
     invalidateRegions();
-    m_needsColumnHeightsRecalculation = heightIsAuto();
+    m_needsColumnHeightsRecalculation = true;
     layout();
 }
 
 bool RenderMultiColumnFlowThread::recalculateColumnHeights()
 {
     // All column sets that needed layout have now been laid out, so we can finally validate them.
     validateRegions();
 
     if (!m_needsColumnHeightsRecalculation)
         return false;
@@ skipping 37 matching lines @@
         width = std::max<LayoutUnit>(0, (availableWidth - ((count - 1) * columnGap)) / count);
     } else if (!columnStyle->hasAutoColumnWidth() && columnStyle->hasAutoColumnCount()) {
         count = std::max<LayoutUnit>(1, (availableWidth + columnGap) / (computedColumnWidth + columnGap));
         width = ((availableWidth + columnGap) / count) - columnGap;
     } else {
         count = std::max<LayoutUnit>(std::min<LayoutUnit>(computedColumnCount, (availableWidth + columnGap) / (computedColumnWidth + columnGap)), 1);
         width = ((availableWidth + columnGap) / count) - columnGap;
     }
 }
 
+bool RenderMultiColumnFlowThread::isDescendantValidColumnSpanner(RenderObject* descendant) const
+{
+    // We assume that we're inside the flow thread. This function is not to be called otherwise.
+    ASSERT(descendant->isDescendantOf(this));
+
+    // First make sure that the renderer itself has the right properties for becoming a spanner.
+    if (descendant->style()->columnSpan() != ColumnSpanAll || !descendant->isBox() || descendant->isFloatingOrOutOfFlowPositioned())
+        return false;
+
+    RenderBlock* container = descendant->containingBlock();
+    if (!container->isRenderBlockFlow() || container->childrenInline()) {
+        // Needs to be block-level.
+        return false;
+    }
+
+    // This looks like a spanner, but if we're inside something unbreakable, it's not to be treated as one.
+    for (RenderBox* ancestor = toRenderBox(descendant)->parentBox(); ancestor; ancestor = ancestor->parentBox()) {
+        if (ancestor->isRenderFlowThread()) {
+            // Don't allow any intervening non-multicol fragmentation contexts. The spec doesn't say
+            // anything about disallowing this, but it's just going to be too complicated to
+            // implement (not to mention specify behavior).
+            return ancestor == this;
+        }
+        if (m_spannerMap.get(ancestor))
+            return false; // Nested spanners not allowed.
+        if (ancestor->isUnsplittableForPagination())
+            return false;
+    }
+    ASSERT_NOT_REACHED();
+    return false;
+}
+
 const char* RenderMultiColumnFlowThread::renderName() const
 {
     return "RenderMultiColumnFlowThread";
 }
 
 void RenderMultiColumnFlowThread::addRegionToThread(RenderMultiColumnSet* columnSet)
 {
     if (RenderMultiColumnSet* nextSet = columnSet->nextSiblingMultiColumnSet()) {
         RenderMultiColumnSetList::iterator it = m_multiColumnSetList.find(nextSet);
         ASSERT(it != m_multiColumnSetList.end());
         m_multiColumnSetList.insertBefore(it, columnSet);
     } else {
         m_multiColumnSetList.add(columnSet);
     }
     columnSet->setIsValid(true);
 }
 
 void RenderMultiColumnFlowThread::willBeRemovedFromTree()
 {
+    m_spannerMap.clear();
     // Detach all column sets from the flow thread. Cannot destroy them at this point, since they
     // are siblings of this object, and there may be pointers to this object's sibling somewhere
     // further up on the call stack.
     for (RenderMultiColumnSet* columnSet = firstMultiColumnSet(); columnSet; columnSet = columnSet->nextSiblingMultiColumnSet())
         columnSet->detachRegion();
     multiColumnBlockFlow()->resetMultiColumnFlowThread();
     RenderFlowThread::willBeRemovedFromTree();
 }
 
+bool RenderMultiColumnFlowThread::isColumnSpanner(const RenderObject* descendant) const
+{
+    ASSERT(descendant->isDescendantOf(this));
+    return m_spannerMap.get(descendant);
+}
+
+bool RenderMultiColumnFlowThread::isInsideColumnSpanner(const RenderObject* descendant) const
+{
+    ASSERT(descendant->isDescendantOf(this));
+    return containingColumnSpanner(descendant);
+}
+
+LayoutUnit RenderMultiColumnFlowThread::enterColumnSpanner(RenderBox* renderer, LayoutUnit logicalTop, SubtreeLayoutScope& layoutScope)
+{
+    ASSERT(renderer->isDescendantOf(this));
+    RenderMultiColumnSpannerSet* spannerSet = m_spannerMap.get(renderer);
+    ASSERT(spannerSet);
+
+    // FIXME: it's really only necessary to mark the spanner set for layout when the height of
+    // |renderer| changes.
+    spannerSet->setChildNeedsLayout(MarkOnlyThis, &layoutScope);
+
+    RenderMultiColumnSet* previousSet = spannerSet->previousSiblingMultiColumnSet();
+    if (!previousSet) {
+        // The first set is entered at the beginning of flow thread layout. If the first set happens
+        // to be a spanner, we have nothing more to do here.
+        return LayoutUnit();
+    }
+
+    RenderBlock* cb = renderer->containingBlock();
+    LayoutUnit logicalTopInFlowThread = cb->offsetFromLogicalTopOfFirstPage() + logicalTop;
+    LayoutUnit adjustment;
+    if (!previousSet->isRenderMultiColumnSpannerSet() && previousSet->pageLogicalHeight()) {
+        // Pad flow thread offset to a column boundary, so that contents that's supposed to come
+        // after the spanner (or the spanner itself) don't bleed into the column preceding the
+        // spanner.
+        LayoutUnit columnLogicalTopInFlowThread = previousSet->pageLogicalTopForOffset(logicalTopInFlowThread);
+        if (columnLogicalTopInFlowThread != logicalTopInFlowThread) {
+            adjustment = columnLogicalTopInFlowThread + previousSet->pageLogicalHeight() - logicalTopInFlowThread;
+            logicalTopInFlowThread += adjustment;
+        }
+    }
+
+    if (!previousSet->isRenderMultiColumnSpannerSet())
+        previousSet->endFlow(logicalTopInFlowThread);
+    spannerSet->beginFlow(logicalTopInFlowThread);
+
+    m_lastSetWorkedOn = spannerSet;
+    return adjustment;
+}
+
+void RenderMultiColumnFlowThread::leaveColumnSpanner(RenderBox* renderer, LayoutUnit logicalBottom)
+{
+    ASSERT(m_lastSetWorkedOn == m_spannerMap.get(renderer));
+
+    RenderBlock* cb = renderer->containingBlock();
+    LayoutUnit logicalBottomInFlowThread = cb->offsetFromLogicalTopOfFirstPage() + logicalBottom;
+    m_lastSetWorkedOn->endFlow(logicalBottomInFlowThread);
+    RenderMultiColumnSet* nextSet = m_lastSetWorkedOn->nextSiblingMultiColumnSet();
+    if (nextSet) {
+        m_lastSetWorkedOn = nextSet;
+        if (!m_lastSetWorkedOn->isRenderMultiColumnSpannerSet())
+            m_lastSetWorkedOn->beginFlow(logicalBottomInFlowThread);
+    }
+}
+
+void RenderMultiColumnFlowThread::flowThreadDescendantInserted(RenderObject* descendant)
+{
+    ASSERT(!m_beingEvacuated);
+    RenderBlockFlow* multicolContainer = multiColumnBlockFlow();
+    for (RenderObject* renderer = descendant; renderer; renderer = renderer->nextInPreOrder(descendant)) {
+        RenderObject* nextRenderer = renderer->nextInPreOrderAfterChildren(this);
+        RenderMultiColumnSet* insertBeforeSet = 0;
+        if (isDescendantValidColumnSpanner(renderer)) {
+            ASSERT(!containingColumnSpanner(renderer));
+            RenderMultiColumnSet* setToSplit = 0;
+            if (nextRenderer) {
+                if (nextRenderer->isColumnSpanAll()) {
+                    insertBeforeSet = m_spannerMap.get(nextRenderer);
+                } else {
+                    RenderObject* previousRenderer = renderer->previousInPreOrder(this);
+                    if (previousRenderer) {
+                        if (RenderMultiColumnSpannerSet* previousSpanner = containingColumnSpanner(previousRenderer)) {
+                            insertBeforeSet = previousSpanner->nextSiblingMultiColumnSet();
+                        } else {
+                            // This is in the middle of an existing column set. Need to split it and put the
+                            // spanner between them.
+                            setToSplit = findSetRendering(previousRenderer);
+                            setToSplit->setNeedsLayoutAndFullPaintInvalidation();
+                            insertBeforeSet = setToSplit->nextSiblingMultiColumnSet();
+                        }
+                    } else {
+                        insertBeforeSet = firstMultiColumnSet();
+                    }
+                }
+            }
+            // Insert the spanner.
+            RenderMultiColumnSpannerSet* newSpanner = RenderMultiColumnSpannerSet::createAnonymous(this, multicolContainer->style(), toRenderBox(renderer));
+            multicolContainer->RenderBlock::addChild(newSpanner, insertBeforeSet);
+            m_spannerMap.add(renderer, newSpanner);
+            invalidateRegions();
+            if (!setToSplit)
+                continue;
+        } else {
+            if (containingColumnSpanner(renderer))
+                continue;
+            if (nextRenderer) {
+                if (RenderMultiColumnSpannerSet* spanner = containingColumnSpanner(nextRenderer)) {
+                    // Inserted right before a spanner. Is there a set for us there?
+                    RenderMultiColumnSet* previous = spanner->previousSiblingMultiColumnSet();
+                    if (previous && !previous->isRenderMultiColumnSpannerSet())
+                        continue;
+                    insertBeforeSet = spanner;
+                } else if (lastMultiColumnSet()) {
+                    continue;
+                }
+            } else if (RenderMultiColumnSet* lastSet = lastMultiColumnSet()) {
+                if (!lastSet->isRenderMultiColumnSpannerSet())
+                    continue;
+            }
+        }
+        RenderMultiColumnSet* newSet = RenderMultiColumnSet::createAnonymous(this, multicolContainer->style());
+        multicolContainer->RenderBlock::addChild(newSet, insertBeforeSet);
+        invalidateRegions();
+
+        // We cannot handle immediate column set siblings at the moment (and there's no need for
+        // it, either). There has to be at least one spanner separating them.
+        ASSERT(!newSet->previousSiblingMultiColumnSet() || newSet->previousSiblingMultiColumnSet()->isRenderMultiColumnSpannerSet());
+        ASSERT(!newSet->nextSiblingMultiColumnSet() || newSet->nextSiblingMultiColumnSet()->isRenderMultiColumnSpannerSet());
+    }
+}
+
+void RenderMultiColumnFlowThread::flowThreadDescendantWillBeRemoved(RenderObject* descendant)
+{
+    if (m_beingEvacuated)
+        return;
+    RenderMultiColumnSpannerSet* containingSpanner = containingColumnSpanner(descendant);
+    if (containingSpanner && containingSpanner->rendererInFlowThread() != descendant)
+        return; // Only removing stuff inside a spanner, and not the spanner itself. Nothing to see here.
+    RenderObject* next;
+    for (RenderObject* renderer = descendant; renderer; renderer = next) {
+        RenderMultiColumnSpannerSet* spanner = m_spannerMap.get(renderer);
+        if (!spanner) {
+            next = renderer->nextInPreOrder(descendant);
+            continue;
+        }
+        next = renderer->nextInPreOrderAfterChildren(descendant); // It's a spanner. Its children are of no interest to us.
+        if (RenderMultiColumnSet* nextSet = spanner->nextSiblingMultiColumnSet()) {
+            RenderMultiColumnSet* previousSet = spanner->previousSiblingMultiColumnSet();
+            if (nextSet && !nextSet->isRenderMultiColumnSpannerSet()
+                && previousSet && !previousSet->isRenderMultiColumnSpannerSet()) {
+                // Need to merge two column sets.
+                nextSet->destroy();
+                previousSet->setNeedsLayout();
+                invalidateRegions();
+            }
+        }
+        m_spannerMap.remove(renderer);
+        spanner->destroy();
+    }
+    if (containingSpanner)
+        return; // Column content has not been removed.
+
+    // Get rid of column sets we no longer need.
+    RenderMultiColumnSpannerSet* adjacentPreviousSpanner = 0;
+    RenderObject* previousRenderer = descendant->previousInPreOrder(this);
+    if (previousRenderer) {
+        adjacentPreviousSpanner = containingColumnSpanner(previousRenderer);
+        if (!adjacentPreviousSpanner)
+            return; // Preceded by column content. Set still needed.
+    }
+    RenderMultiColumnSpannerSet* adjacentNextSpanner = 0;
+    RenderObject* nextRenderer = descendant->nextInPreOrderAfterChildren(this);
+    if (nextRenderer) {
+        adjacentNextSpanner = containingColumnSpanner(nextRenderer);
+        if (!adjacentNextSpanner)
+            return; // Followed by column content. Set still needed.
+    }
+    RenderMultiColumnSet* columnSetToRemove;
+    if (adjacentNextSpanner) {
+        columnSetToRemove = adjacentNextSpanner->previousSiblingMultiColumnSet();
+    } else if (adjacentPreviousSpanner) {
+        columnSetToRemove = adjacentPreviousSpanner->nextSiblingMultiColumnSet();
+    } else {
+        columnSetToRemove = firstMultiColumnSet();
+        ASSERT(columnSetToRemove);
+        ASSERT(!columnSetToRemove->previousSiblingMultiColumnSet());
+        ASSERT(!columnSetToRemove->nextSiblingMultiColumnSet());
+    }
+    ASSERT(columnSetToRemove);
+    columnSetToRemove->destroy();
+}
+
+void RenderMultiColumnFlowThread::flowThreadDescendantStyleDidChange(RenderObject* descendant)
+{
+    ASSERT(descendant->isDescendantOf(this));
+    if (RenderMultiColumnSpannerSet* spanner = m_spannerMap.get(descendant))
+        spanner->updateMarginProperties();
+}
+
 void RenderMultiColumnFlowThread::computeLogicalHeight(LayoutUnit logicalHeight, LayoutUnit logicalTop, LogicalExtentComputedValues& computedValues) const
 {
     // We simply remain at our intrinsic height.
     computedValues.m_extent = logicalHeight;
     computedValues.m_position = logicalTop;
 }
 
 void RenderMultiColumnFlowThread::updateLogicalWidth()
 {
     LayoutUnit columnWidth;
     calculateColumnCountAndWidth(columnWidth, m_columnCount);
     setLogicalWidth(columnWidth);
 }
 
 void RenderMultiColumnFlowThread::layout()
 {
+    ASSERT(!m_lastSetWorkedOn);
+    m_lastSetWorkedOn = firstMultiColumnSet();
+    if (m_lastSetWorkedOn)
+        m_lastSetWorkedOn->beginFlow(LayoutUnit());
     RenderFlowThread::layout();
-    if (RenderMultiColumnSet* lastSet = lastMultiColumnSet())
+    if (RenderMultiColumnSet* lastSet = lastMultiColumnSet()) {
+        lastSet->endFlow(logicalHeight());
         lastSet->expandToEncompassFlowThreadContentsIfNeeded();
+    }
+    m_lastSetWorkedOn = 0;
 }
 
 void RenderMultiColumnFlowThread::setPageBreak(LayoutUnit offset, LayoutUnit spaceShortage)
 {
     // Only positive values are interesting (and allowed) here. Zero space shortage may be reported
     // when we're at the top of a column and the element has zero height. Ignore this, and also
     // ignore any negative values, which may occur when we set an early break in order to honor
     // widows in the next column.
     if (spaceShortage <= 0)
         return;
 
     if (RenderMultiColumnSet* multicolSet = columnSetAtBlockOffset(offset))
         multicolSet->recordSpaceShortage(spaceShortage);
 }
 
 void RenderMultiColumnFlowThread::updateMinimumPageHeight(LayoutUnit offset, LayoutUnit minHeight)
 {
     if (RenderMultiColumnSet* multicolSet = columnSetAtBlockOffset(offset))
         multicolSet->updateMinimumColumnHeight(minHeight);
 }
 
-RenderMultiColumnSet* RenderMultiColumnFlowThread::columnSetAtBlockOffset(LayoutUnit /*offset*/) const
+RenderMultiColumnSet* RenderMultiColumnFlowThread::columnSetAtBlockOffset(LayoutUnit offset) const
 {
-    // For now there's only one column set, so this is easy:
-    return firstMultiColumnSet();
+    if (m_lastSetWorkedOn) {
+        // Layout in progress. We are calculating the set heights as we speak, so the column set range
+        // information is not up-to-date.
+        RenderMultiColumnSet* columnSet = m_lastSetWorkedOn;
+        if (offset < columnSet->logicalTopInFlowThread()) {
+            // In some cases we need to search backwards for a column set we've advanced past. This
+            // happens when a block crosses a column set boundary, and someone wants to examine or
+            // adjust its top after or during layout. FIXME: If its top acually gets adjusted
+            // (e.g. because of an incorrect initial top position estimate), this may be problematic
+            // for column balancing, but returning the correct set here is at least better than
+            // nothing.
+            do {
+                if (RenderMultiColumnSet* prev = columnSet->previousSiblingMultiColumnSet())
+                    columnSet = prev;
+                else
+                    break;
+            } while (offset < columnSet->logicalTopInFlowThread());
+        }
+        return columnSet;
+    }
+
+    ASSERT(!m_regionsInvalidated);
+    if (offset <= 0)
+        return m_multiColumnSetList.isEmpty() ? 0 : m_multiColumnSetList.first();
+
+    MultiColumnSetSearchAdapter adapter(offset);
+    m_multiColumnSetIntervalTree.allOverlapsWithAdapter<MultiColumnSetSearchAdapter>(adapter);
+
+    // If no set was found, the offset is in the flow thread overflow.
+    if (!adapter.result() && !m_multiColumnSetList.isEmpty())
+        return m_multiColumnSetList.last();
+    return adapter.result();
 }
 
 bool RenderMultiColumnFlowThread::addForcedRegionBreak(LayoutUnit offset, RenderObject* /*breakChild*/, bool /*isBefore*/, LayoutUnit* offsetBreakAdjustment)
 {
     if (RenderMultiColumnSet* multicolSet = columnSetAtBlockOffset(offset)) {
         multicolSet->addContentRun(offset);
         if (offsetBreakAdjustment)
             *offsetBreakAdjustment = pageLogicalHeightForOffset(offset) ? pageRemainingLogicalHeightForOffset(offset, IncludePageBoundary) : LayoutUnit();
         return true;
     }
     return false;
 }
 
 bool RenderMultiColumnFlowThread::isPageLogicalHeightKnown() const
 {
     if (RenderMultiColumnSet* columnSet = lastMultiColumnSet())
         return columnSet->pageLogicalHeight();
     return false;
 }
 
 }