Fixed flexbox layout bug when container size changes

Added cross-axis alignment functionality

src/elements/Flexbox.lua

@@ -31,6 +31,19 @@ Flexbox.defineProperty(Flexbox, "flexJustifyContent", {
         return value
     end
 })
+---@property flexAlignItems string "flex-start" The alignment of flex items along the cross axis
+Flexbox.defineProperty(Flexbox, "flexAlignItems", {
+    default = "flex-start",
+    type = "string",
+    setter = function(self, value)
+        if not value:match("^flex%-") and value ~= "stretch" then
+            value = "flex-" .. value
+        end
+        return value
+    end
+})
+---@property flexCrossPadding number 0 The padding on both sides of the cross axis
+Flexbox.defineProperty(Flexbox, "flexCrossPadding", {default = 0, type = "number"})
 ---@property flexWrap boolean false Whether to wrap flex items onto multiple lines
 ---@property flexUpdateLayout boolean false Whether to update the layout of the flexbox
 Flexbox.defineProperty(Flexbox, "flexWrap", {default = false, type = "boolean"})
@@ -57,187 +70,576 @@ local lineBreakElement = {
 
 
 local function sortElements(self, direction, spacing, wrap)
+    -- Pre-allocate tables to reduce dynamic expansion
     local elements = self.get("children")
     local sortedElements = {}
-    if not(wrap)then
+    local visibleElements = {}
-      local index = 1
+    local childCount = 0
-      local lineSize = 1
+    
-      local lineOffset = 1
+    -- First calculate the number of visible elements, pre-allocate space
-      for _,v in pairs(elements)do
+    for _, elem in pairs(elements) do
-          if(sortedElements[index]==nil)then sortedElements[index]={offset=1} end
+        if elem ~= lineBreakElement and elem:getVisible() then
-
+            childCount = childCount + 1
-          local childHeight = direction == "row" and v.get("height") or v.get("width")
+        end
-          if childHeight > lineSize then
+    end
-              lineSize = childHeight
+    
-          end
+    -- Use known size to pre-allocate array
-          if(v == lineBreakElement)then
+    if not wrap then
-              lineOffset = lineOffset + lineSize + spacing
+        -- No-wrap mode, all elements in one row/column
-              lineSize = 1
+        sortedElements[1] = {offset=1}
-              index = index + 1
+        
-              sortedElements[index] = {offset=lineOffset}
+        for _, elem in pairs(elements) do
-          else
+            if elem == lineBreakElement then
-              table.insert(sortedElements[index], v)
+                -- Create new line
-          end
+                local nextIndex = #sortedElements + 1
-      end
+                if sortedElements[nextIndex] == nil then
-    elseif(wrap)then
+                    sortedElements[nextIndex] = {offset=1}
-      local lineSize = 1
+                end
-      local lineOffset = 1
+            else
-
+                table.insert(sortedElements[#sortedElements], elem)
-      local maxSize = direction == "row" and self.get("width") or self.get("height")
+            end
-      local usedSize = 0
+        end
-      local index = 1
+    else
-
+        -- Wrap mode, need to calculate rows/columns
-      for _,v in pairs(elements) do
+        local index = 1
-          if(sortedElements[index]==nil) then sortedElements[index]={offset=1} end
+        local containerSize = direction == "row" and self.get("width") or self.get("height")
-
+        local remainingSpace = containerSize
-          if v:getType() == "lineBreak" then
+        sortedElements[index] = {offset=1}
-              lineOffset = lineOffset + lineSize + spacing
+        
-              usedSize = 0
+        for _, elem in pairs(elements) do
-              lineSize = 1
+            if elem == lineBreakElement then
-              index = index + 1
+                -- Create new line
-              sortedElements[index] = {offset=lineOffset}
+                index = index + 1
-          else
+                sortedElements[index] = {offset=1}
-              local objSize = direction == "row" and v.get("width") or v.get("height")
+                remainingSpace = containerSize
-              if(objSize+usedSize<=maxSize) then
+            else
-                  table.insert(sortedElements[index], v)
+                local elemSize = direction == "row" and elem.get("width") or elem.get("height")
-                  usedSize = usedSize + objSize + spacing
+                if elemSize + spacing <= remainingSpace then
-              else
+                    -- Element fits in current line
-                  lineOffset = lineOffset + lineSize + spacing
+                    table.insert(sortedElements[index], elem)
-                  lineSize = direction == "row" and v.get("height") or v.get("width")
+                    remainingSpace = remainingSpace - elemSize - spacing
-                  index = index + 1
+                else
-                  usedSize = objSize + spacing
+                    -- Need new line
-                  sortedElements[index] = {offset=lineOffset, v}
+                    index = index + 1
-              end
+                    sortedElements[index] = {offset=1, elem}
-
+                    remainingSpace = containerSize - elemSize - spacing
-              local childHeight = direction == "row" and v.get("height") or v.get("width")
+                end
-              if childHeight > lineSize then
+            end
-                  lineSize = childHeight
+        end
-              end
+    end
-          end
+    
-      end
-  end
     return sortedElements
 end
 
 local function calculateRow(self, children, spacing, justifyContent)
     local containerWidth = self.get("width")
-
+    local containerHeight = self.get("height")
-    local usedSpace = spacing * (#children - 1)
+    local alignItems = self.get("flexAlignItems")
+    local crossPadding = self.get("flexCrossPadding")
+    
+    -- Safety check
+    if containerWidth <= 0 then return end
+    
+    -- Calculate available cross axis space (considering padding)
+    local availableCrossAxisSpace = containerHeight - (crossPadding * 2)
+    if availableCrossAxisSpace < 1 then
+        availableCrossAxisSpace = containerHeight
+        crossPadding = 0
+    end
+    
+    -- Cache local variables to reduce function calls
+    local max = math.max
+    local min = math.min
+    local floor = math.floor
+    local ceil = math.ceil
+    
+    -- Fixed elements
+    local fixedElements = {}
+    -- Flexible elements
+    local flexElements = {}
+    -- Total flex coefficient
     local totalFlexGrow = 0
-
+    -- Pre-allocate capacity
+    local fixedCount = 0
+    local flexCount = 0
+    
+    -- First calculate element counts to pre-allocate space
     for _, child in ipairs(children) do
         if child ~= lineBreakElement then
-            usedSpace = usedSpace + child.get("width")
+            local grow = child.get("flexGrow") or 0
-            totalFlexGrow = totalFlexGrow + child.get("flexGrow")
+            if grow > 0 then
+                flexCount = flexCount + 1
+            else
+                fixedCount = fixedCount + 1
+            end
         end
     end
-
+    
-    local remainingSpace = containerWidth - usedSpace
+    -- Pre-allocate table space
-    local extraSpacePerUnit = totalFlexGrow > 0 and (remainingSpace / totalFlexGrow) or 0
+    for i = 1, fixedCount do fixedElements[i] = nil end
-    local distributedSpace = 0
+    
-
+    -- Step 1: Categorize elements and collect information
-    local currentX = 1
+    for _, child in ipairs(children) do
-    for i, child in ipairs(children) do
         if child ~= lineBreakElement then
-            local childWidth = child.get("width")
+            local grow = child.get("flexGrow") or 0
-
+            if grow > 0 then
-            if child.get("flexGrow") > 0 then
+                totalFlexGrow = totalFlexGrow + grow
-
+                table.insert(flexElements, {element = child, grow = grow})
-                if i == #children then
+            else
-                    local extraSpace = remainingSpace - distributedSpace
+                table.insert(fixedElements, child)
-                    childWidth = childWidth + extraSpace
+            end
-                else
+        end
-                    local extraSpace = math.floor(extraSpacePerUnit * child.get("flexGrow"))
+    end
-                    childWidth = childWidth + extraSpace
+    
-                    distributedSpace = distributedSpace + extraSpace
+    -- Step 2: Pre-processing before layout
+    -- First calculate the total width needed for all fixed elements
+    local fixedWidthSum = 0
+    for _, element in ipairs(fixedElements) do
+        fixedWidthSum = fixedWidthSum + element.get("width")
+    end
+    
+    -- Calculate total width of gaps
+    local totalElements = #fixedElements + #flexElements
+    local totalGaps = totalElements > 1 and (totalElements - 1) or 0
+    local gapsWidth = spacing * totalGaps
+    
+    -- Calculate total available space for flexible elements
+    local flexAvailableSpace = max(0, containerWidth - fixedWidthSum - gapsWidth)
+    
+    -- Safety check: If not enough space, force compress fixed elements
+    if flexAvailableSpace < 0 then
+        -- Set gaps to zero
+        gapsWidth = 0
+        flexAvailableSpace = containerWidth - fixedWidthSum
+        
+        -- If still not enough, need to shrink fixed elements
+        if flexAvailableSpace < 0 and #fixedElements > 0 then
+            local reductionPerElement = ceil(-flexAvailableSpace / #fixedElements)
+            for _, element in ipairs(fixedElements) do
+                local currentWidth = element.get("width")
+                local newWidth = max(1, currentWidth - reductionPerElement)
+                element.set("width", newWidth)
+                flexAvailableSpace = flexAvailableSpace + (currentWidth - newWidth)
+                if flexAvailableSpace >= 0 then
+                    break
                 end
             end
-
+        end
-            child.set("x", currentX)
+        
-            child.set("y", children.offset or 1)
+        -- If still not enough, may need to set minimum width
-            child.set("width", childWidth)
+        flexAvailableSpace = max(0, flexAvailableSpace)
-            currentX = currentX + childWidth + spacing
+    end
+    
+    -- Step 3: Allocate space for flexible elements
+    -- Pre-allocate table to avoid dynamic expansion
+    local allocatedWidths = {}
+    for i = 1, flexCount do
+        allocatedWidths[flexElements[i].element] = nil
+    end
+    
+    -- If there are flexible elements and available space
+    if #flexElements > 0 and flexAvailableSpace > 0 and totalFlexGrow > 0 then
+        -- Reserve some safety margin (e.g., 5% of space) to ensure no overflow due to rounding
+        local safeFlexSpace = floor(flexAvailableSpace * 0.95)
+        
+        -- Allocate base width for each element (conservative strategy)
+        for _, item in ipairs(flexElements) do
+            -- Determine this element's share
+            local proportion = item.grow / totalFlexGrow
+            -- Determine width to allocate (floor to ensure safety)
+            local extraWidth = floor(safeFlexSpace * proportion)
+            -- Set final width
+            allocatedWidths[item.element] = item.element.get("width") + extraWidth
         end
     end
-
+    
-    if justifyContent == "flex-end" then
+    -- Step 4: Strictly validate final widths
-        local offset = containerWidth - (currentX - spacing - 1)
+    -- Calculate total width after allocation (including gaps)
-        for _, child in ipairs(children) do
+    local finalTotalWidth = gapsWidth
-            child.set("x", child.get("x") + offset)
+    for _, element in ipairs(fixedElements) do
+        finalTotalWidth = finalTotalWidth + element.get("width")
+    end
+    for _, item in ipairs(flexElements) do
+        local width = allocatedWidths[item.element] or item.element.get("width")
+        finalTotalWidth = finalTotalWidth + width
+    end
+    
+    -- If total width exceeds container, proportionally reduce all elements
+    if finalTotalWidth > containerWidth then
+        local excessWidth = finalTotalWidth - containerWidth
+        local reductionFactor = excessWidth / (finalTotalWidth - gapsWidth)
+        
+        -- First reduce flexible elements
+        if #flexElements > 0 then
+            for _, item in ipairs(flexElements) do
+                local width = allocatedWidths[item.element] or item.element.get("width")
+                local reduction = ceil(width * reductionFactor)
+                allocatedWidths[item.element] = max(1, width - reduction)
+            end
         end
-    elseif justifyContent == "flex-center" or justifyContent == "center" then -- Akzeptiere beide Formate
+        
-        local offset = math.floor((containerWidth - (currentX - spacing - 1)) / 2)
+        -- If still not enough, reduce fixed elements
-        for _, child in ipairs(children) do
+        finalTotalWidth = gapsWidth
-            child.set("x", child.get("x") + offset)
+        for _, element in ipairs(fixedElements) do
+            finalTotalWidth = finalTotalWidth + element.get("width")
+        end
+        for _, item in ipairs(flexElements) do
+            finalTotalWidth = finalTotalWidth + (allocatedWidths[item.element] or item.element.get("width"))
+        end
+        
+        if finalTotalWidth > containerWidth and #fixedElements > 0 then
+            excessWidth = finalTotalWidth - containerWidth
+            reductionFactor = excessWidth / (finalTotalWidth - gapsWidth)
+            
+            for _, element in ipairs(fixedElements) do
+                local width = element.get("width")
+                local reduction = ceil(width * reductionFactor)
+                element.set("width", max(1, width - reduction))
+            end
+        end
+    end
+    
+    -- Step 5: Apply layout
+    local currentX = 1
+    
+    -- Place all elements
+    for _, child in ipairs(children) do
+        if child ~= lineBreakElement then
+            -- Apply X coordinate
+            child.set("x", currentX)
+            
+            -- Apply Y coordinate (based on vertical alignment)
+            if alignItems == "stretch" then
+                -- Vertical stretch to fill container, considering padding
+                child.set("height", availableCrossAxisSpace)
+                child.set("y", 1 + crossPadding)
+            else
+                local childHeight = child.get("height")
+                local y = 1
+                
+                if alignItems == "flex-end" then
+                    -- Bottom align
+                    y = containerHeight - childHeight + 1
+                elseif alignItems == "flex-center" or alignItems == "center" then
+                    -- Center align
+                    y = floor((containerHeight - childHeight) / 2) + 1
+                end
+                
+                -- Ensure Y value is not less than 1
+                child.set("y", max(1, y))
+            end
+            
+            -- If flexible element, apply calculated width
+            if allocatedWidths[child] then
+                child.set("width", allocatedWidths[child])
+            end
+            
+            -- Final safety check (using cached math functions)
+            local rightEdge = currentX + child.get("width") - 1
+            if rightEdge > containerWidth then
+                child.set("width", max(1, containerWidth - currentX + 1))
+            end
+            
+            -- Final safety check height doesn't exceed container
+            local bottomEdge = child.get("y") + child.get("height") - 1
+            if bottomEdge > containerHeight then
+                child.set("height", max(1, containerHeight - child.get("y") + 1))
+            end
+            
+            -- Update position for next element
+            currentX = currentX + child.get("width") + spacing
+            
+            -- Ensure won't exceed container right edge
+            if currentX > containerWidth + 1 then
+                currentX = containerWidth + 1
+            end
+        end
+    end
+    
+    -- Apply alignment (only when remaining space is positive)
+    local usedWidth = min(containerWidth, currentX - spacing - 1)
+    local remainingSpace = containerWidth - usedWidth
+    
+    if remainingSpace > 0 then
+        if justifyContent == "flex-end" then
+            for _, child in ipairs(children) do
+                if child ~= lineBreakElement then
+                    child.set("x", child.get("x") + remainingSpace)
+                end
+            end
+        elseif justifyContent == "flex-center" or justifyContent == "center" then
+            local offset = floor(remainingSpace / 2)
+            for _, child in ipairs(children) do
+                if child ~= lineBreakElement then
+                    child.set("x", child.get("x") + offset)
+                end
+            end
         end
     end
 end
 
 local function calculateColumn(self, children, spacing, justifyContent)
+    local containerWidth = self.get("width")
     local containerHeight = self.get("height")
-
+    local alignItems = self.get("flexAlignItems")
-    local usedSpace = spacing * (#children - 1)
+    local crossPadding = self.get("flexCrossPadding")
+    
+    -- Safety check
+    if containerHeight <= 0 then return end
+    
+    -- Calculate available cross axis space (considering padding)
+    local availableCrossAxisSpace = containerWidth - (crossPadding * 2)
+    if availableCrossAxisSpace < 1 then
+        availableCrossAxisSpace = containerWidth
+        crossPadding = 0
+    end
+    
+    -- Cache local variables to reduce function calls
+    local max = math.max
+    local min = math.min
+    local floor = math.floor
+    local ceil = math.ceil
+    
+    -- Fixed elements
+    local fixedElements = {}
+    -- Flexible elements
+    local flexElements = {}
+    -- Total flex coefficient
     local totalFlexGrow = 0
-
+    -- Pre-allocate capacity
+    local fixedCount = 0
+    local flexCount = 0
+    
+    -- First calculate element counts to pre-allocate space
     for _, child in ipairs(children) do
         if child ~= lineBreakElement then
-            usedSpace = usedSpace + child.get("height")
+            local grow = child.get("flexGrow") or 0
-            totalFlexGrow = totalFlexGrow + child.get("flexGrow")
+            if grow > 0 then
+                flexCount = flexCount + 1
+            else
+                fixedCount = fixedCount + 1
+            end
         end
     end
-
+    
-    local remainingSpace = containerHeight - usedSpace
+    -- Pre-allocate table space
-    local extraSpacePerUnit = totalFlexGrow > 0 and (remainingSpace / totalFlexGrow) or 0
+    for i = 1, fixedCount do fixedElements[i] = nil end
-    local distributedSpace = 0
+    
-
+    -- Step 1: Categorize elements and collect information
-    local currentY = 1
+    for _, child in ipairs(children) do
-    for i, child in ipairs(children) do
         if child ~= lineBreakElement then
-            local childHeight = child.get("height")
+            local grow = child.get("flexGrow") or 0
-
+            if grow > 0 then
-            if child.get("flexGrow") > 0 then
+                totalFlexGrow = totalFlexGrow + grow
-
+                table.insert(flexElements, {element = child, grow = grow})
-                if i == #children then
+            else
-                    local extraSpace = remainingSpace - distributedSpace
+                table.insert(fixedElements, child)
-                    childHeight = childHeight + extraSpace
+            end
-                else
+        end
-                    local extraSpace = math.floor(extraSpacePerUnit * child.get("flexGrow"))
+    end
-                    childHeight = childHeight + extraSpace
+    
-                    distributedSpace = distributedSpace + extraSpace
+    -- Step 2: Pre-processing before layout
+    
+    -- First calculate the total height needed for all fixed elements
+    local fixedHeightSum = 0
+    for _, element in ipairs(fixedElements) do
+        fixedHeightSum = fixedHeightSum + element.get("height")
+    end
+    
+    -- Calculate total height of gaps
+    local totalElements = #fixedElements + #flexElements
+    local totalGaps = totalElements > 1 and (totalElements - 1) or 0
+    local gapsHeight = spacing * totalGaps
+    
+    -- Calculate total available space for flexible elements
+    local flexAvailableSpace = max(0, containerHeight - fixedHeightSum - gapsHeight)
+    
+    -- Safety check: If not enough space, force compress fixed elements
+    if flexAvailableSpace < 0 then
+        -- Set gaps to zero
+        gapsHeight = 0
+        flexAvailableSpace = containerHeight - fixedHeightSum
+        
+        -- If still not enough, need to shrink fixed elements
+        if flexAvailableSpace < 0 and #fixedElements > 0 then
+            local reductionPerElement = ceil(-flexAvailableSpace / #fixedElements)
+            for _, element in ipairs(fixedElements) do
+                local currentHeight = element.get("height")
+                local newHeight = max(1, currentHeight - reductionPerElement)
+                element.set("height", newHeight)
+                flexAvailableSpace = flexAvailableSpace + (currentHeight - newHeight)
+                if flexAvailableSpace >= 0 then
+                    break
                 end
             end
-
+        end
-            child.set("x", children.offset or 1)
+        
-            child.set("y", currentY)
+        -- If still not enough, may need to set minimum height
-            child.set("height", childHeight)
+        flexAvailableSpace = max(0, flexAvailableSpace)
-            currentY = currentY + childHeight + spacing
+    end
+    
+    -- Step 3: Allocate space for flexible elements
+    -- Pre-allocate table to avoid dynamic expansion
+    local allocatedHeights = {}
+    for i = 1, flexCount do
+        allocatedHeights[flexElements[i].element] = nil
+    end
+    
+    -- If there are flexible elements and available space
+    if #flexElements > 0 and flexAvailableSpace > 0 and totalFlexGrow > 0 then
+        -- Reserve some safety margin (e.g., 5% of space) to ensure no overflow due to rounding
+        local safeFlexSpace = floor(flexAvailableSpace * 0.95)
+        
+        -- Allocate base height for each element (conservative strategy)
+        for _, item in ipairs(flexElements) do
+            -- Determine this element's share
+            local proportion = item.grow / totalFlexGrow
+            -- Determine height to allocate (floor to ensure safety)
+            local extraHeight = floor(safeFlexSpace * proportion)
+            -- Set final height
+            allocatedHeights[item.element] = item.element.get("height") + extraHeight
         end
     end
-
+    
-    if justifyContent == "flex-end" then
+    -- Step 4: Strictly validate final heights
-        local offset = containerHeight - (currentY - spacing - 1)
+    -- Calculate total height after allocation (including gaps)
-        for _, child in ipairs(children) do
+    local finalTotalHeight = gapsHeight
-            child.set("y", child.get("y") + offset)
+    for _, element in ipairs(fixedElements) do
+        finalTotalHeight = finalTotalHeight + element.get("height")
+    end
+    for _, item in ipairs(flexElements) do
+        local height = allocatedHeights[item.element] or item.element.get("height")
+        finalTotalHeight = finalTotalHeight + height
+    end
+    
+    -- If total height exceeds container, proportionally reduce all elements
+    if finalTotalHeight > containerHeight then
+        local excessHeight = finalTotalHeight - containerHeight
+        local reductionFactor = excessHeight / (finalTotalHeight - gapsHeight)
+        
+        -- First reduce flexible elements
+        if #flexElements > 0 then
+            for _, item in ipairs(flexElements) do
+                local height = allocatedHeights[item.element] or item.element.get("height")
+                local reduction = ceil(height * reductionFactor)
+                allocatedHeights[item.element] = max(1, height - reduction)
+            end
         end
-    elseif justifyContent == "flex-center" or justifyContent == "center" then -- Akzeptiere beide Formate
+        
-        local offset = math.floor((containerHeight - (currentY - spacing - 1)) / 2)
+        -- If still not enough, reduce fixed elements
-        for _, child in ipairs(children) do
+        finalTotalHeight = gapsHeight
-            child.set("y", child.get("y") + offset)
+        for _, element in ipairs(fixedElements) do
+            finalTotalHeight = finalTotalHeight + element.get("height")
+        end
+        for _, item in ipairs(flexElements) do
+            finalTotalHeight = finalTotalHeight + (allocatedHeights[item.element] or item.element.get("height"))
+        end
+        
+        if finalTotalHeight > containerHeight and #fixedElements > 0 then
+            excessHeight = finalTotalHeight - containerHeight
+            reductionFactor = excessHeight / (finalTotalHeight - gapsHeight)
+            
+            for _, element in ipairs(fixedElements) do
+                local height = element.get("height")
+                local reduction = ceil(height * reductionFactor)
+                element.set("height", max(1, height - reduction))
+            end
+        end
+    end
+    
+    -- Step 5: Apply layout
+    local currentY = 1
+    
+    -- Place all elements
+    for _, child in ipairs(children) do
+        if child ~= lineBreakElement then
+            -- Apply Y coordinate
+            child.set("y", currentY)
+            
+            -- Apply X coordinate (based on horizontal alignment)
+            if alignItems == "stretch" then
+                -- Horizontal stretch to fill container, considering padding
+                child.set("width", availableCrossAxisSpace)
+                child.set("x", 1 + crossPadding)
+            else
+                local childWidth = child.get("width")
+                local x = 1
+                
+                if alignItems == "flex-end" then
+                    -- Right align
+                    x = containerWidth - childWidth + 1
+                elseif alignItems == "flex-center" or alignItems == "center" then
+                    -- Center align
+                    x = floor((containerWidth - childWidth) / 2) + 1
+                end
+                
+                -- Ensure X value is not less than 1
+                child.set("x", max(1, x))
+            end
+            
+            -- If flexible element, apply calculated height
+            if allocatedHeights[child] then
+                child.set("height", allocatedHeights[child])
+            end
+            
+            -- Final safety check (using cached math functions)
+            local bottomEdge = currentY + child.get("height") - 1
+            if bottomEdge > containerHeight then
+                child.set("height", max(1, containerHeight - currentY + 1))
+            end
+            
+            -- Final safety check width doesn't exceed container
+            local rightEdge = child.get("x") + child.get("width") - 1
+            if rightEdge > containerWidth then
+                child.set("width", max(1, containerWidth - child.get("x") + 1))
+            end
+            
+            -- Update position for next element
+            currentY = currentY + child.get("height") + spacing
+            
+            -- Ensure won't exceed container bottom edge
+            if currentY > containerHeight + 1 then
+                currentY = containerHeight + 1
+            end
+        end
+    end
+    
+    -- Apply alignment (only when remaining space is positive)
+    local usedHeight = min(containerHeight, currentY - spacing - 1)
+    local remainingSpace = containerHeight - usedHeight
+    
+    if remainingSpace > 0 then
+        if justifyContent == "flex-end" then
+            for _, child in ipairs(children) do
+                if child ~= lineBreakElement then
+                    child.set("y", child.get("y") + remainingSpace)
+                end
+            end
+        elseif justifyContent == "flex-center" or justifyContent == "center" then
+            local offset = floor(remainingSpace / 2)
+            for _, child in ipairs(children) do
+                if child ~= lineBreakElement then
+                    child.set("y", child.get("y") + offset)
+                end
+            end
         end
     end
 end
 
+-- Optimize updateLayout function
 local function updateLayout(self, direction, spacing, justifyContent, wrap)
-  local elements = sortElements(self, direction, spacing, wrap)
+    -- Get all elements that need layout
-    if direction == "row" then
+    local elements = sortElements(self, direction, spacing, wrap)
-        for _,v in pairs(elements)do
+    
-            calculateRow(self, v, spacing, justifyContent)
+    -- Based on direction, select layout function, avoid checking every iteration
-        end
+    local layoutFunction = direction == "row" and calculateRow or calculateColumn
-    else
+    
-        for _,v in pairs(elements)do
+    -- Apply layout calculation
-            calculateColumn(self, v, spacing, justifyContent)
+    for _, rowOrColumn in pairs(elements) do
-        end
+        layoutFunction(self, rowOrColumn, spacing, justifyContent)
     end
+    
+    -- Reset layout update flag
     self.set("flexUpdateLayout", false)
 end