Search code examples
swiftswift-playground

Playground 401k Calculator: Unable to get "back-door-Roth" function to work properly. Seeking techniques to improve formulae

I'm using Playground as a testing ground for code with the intention to transition in the near future to creating an iOS app.

Background:

In case the term "back-door-Roth" is unfamiliar, let me attempt to level the bubble. If, in a given calendar year, an employee contributes to their 401k up to an amount totaling the IRS contribution limit (currently $22,500 in most cases), the employee typically is restricted by law to discontinue contributions. However, an employer may offer a "back-door" option which would allow the employee to continue contributing beyond the IRS contribution limit. However those "above-the-limit" contributions are placed into a separate Roth account, often labeled a "401a" account. This employee may continue to contribute to their 401a account throughout that calendar year so long as the sum total of his/her contributions and the company's contributions do not exceed the IRS combined contribution limit (currently $66,000 in most cases).

Thanks to Chip Jarred and other SO contributors, the below playground code is fully functional. This code does not allow for 401a "back-door-Roth" contributions, but it is a solid foundation upon which I am building:

    import Foundation
    
    let pennyRoundingBehavior = NSDecimalNumberHandler(
        roundingMode: .bankers,
        scale: 2,
        raiseOnExactness: false,
        raiseOnOverflow: true,
        raiseOnUnderflow: true,
        raiseOnDivideByZero: true
    )
    
    func roundToNearestPenny(percentage: Decimal, of dollarAmount: Decimal) -> Decimal
    {
        assert((0...100).contains(percentage))
        
        let x = ((dollarAmount * percentage / 100) as NSDecimalNumber)
        return x.rounding(accordingToBehavior: pennyRoundingBehavior) as Decimal
    }
    
    enum ContributionMonth: Int, CaseIterable
    {
        case January, February, March, April, May, June,
             July, August, September, October, November, December
        
        static var count:Int  { allCases.count }
    }
    
    struct MonthlyContributions
    {
        var roth                : Decimal = 0
        var traditional         : Decimal = 0
        var rothCompany         : Decimal = 0
        var traditionalCompany  : Decimal = 0
        var yearToDate          : Decimal = 0
        var yearToDateAll       : Decimal = 0
    }
    
    extension MonthlyContributions
    {
        var total: Decimal { roth + traditional }
        var companyTotal: Decimal { rothCompany + traditionalCompany }
        
        func nextMonth(
            combinedAmount: Decimal,
            rothPercentage: Decimal,
            rothPercentageCompany: Decimal,
            combinedAmountCompany: Decimal
            ) -> Self
        {
            let roth = roundToNearestPenny(
                percentage: rothPercentage,
                of: combinedAmount
            )
            let rothCompany = roundToNearestPenny(
                percentage: rothPercentageCompany,
                of: combinedAmountCompany
            )
            return Self(
                roth: roth,
                traditional: combinedAmount - roth,
                rothCompany: rothCompany,
                traditionalCompany: combinedAmountCompany - rothCompany,
                yearToDate: combinedAmount + yearToDate,
                yearToDateAll: combinedAmount + combinedAmountCompany + yearToDateAll
            )
        }
    }
    
    extension Array where Element == MonthlyContributions {
        subscript (index: ContributionMonth) -> Element { self[index.rawValue] }
    }
    
    func calculatePersCombinedContribution(
        monthlyContribution        monthly    : Decimal,
        annualContributionLimit    limit      : Decimal,
        contributionsSoFarThisYear cummulative: Decimal) -> Decimal
    {
        min(max(0, limit - cummulative), monthly)
    }
    func calculateAllCombinedContribution(
        monthlyContribution             monthly : Decimal,
        annualCombinedContributionLimit limit   : Decimal,
        allContributionsSoFarThisYear   cummulative: Decimal) -> Decimal
    {
        min(max(0, limit - cummulative), monthly)
    }
    
    func computeMonthlyContributions(
        monthlyPay            : Decimal,
        contributionPercentage: Decimal,
        companyContributionPercentage: Decimal,
        rothPercentage        : Decimal,
        contributionLimit     : Decimal,
        combinedContributionLimit: Decimal) -> [MonthlyContributions]
    {
        assert((0...100).contains(contributionPercentage))
        assert((0...100).contains(rothPercentage))
        assert(monthlyPay >= 0)
        assert(contributionLimit >= 0)
    
        var contribution = MonthlyContributions()
    
        var monthlyContributions: [MonthlyContributions] = []
        monthlyContributions.reserveCapacity(ContributionMonth.count)
    
        let monthlyContribution = roundToNearestPenny(
            percentage: contributionPercentage,
            of: monthlyPay
        )
        let monthlyCompanyContribution = roundToNearestPenny(
            percentage: companyContributionPercentage,
            of: monthlyPay
        )
        for _ in ContributionMonth.allCases
        {
            let combinedPersAmount = calculatePersCombinedContribution(
                monthlyContribution       : monthlyContribution,
                annualContributionLimit   : contributionLimit,
                contributionsSoFarThisYear: contribution.yearToDate
            )
            let combinedAllAmount = calculateAllCombinedContribution(
                monthlyContribution             : monthlyCompanyContribution,
                annualCombinedContributionLimit : combinedContributionLimit,
                allContributionsSoFarThisYear   : contribution.yearToDateAll
            )
            contribution = contribution.nextMonth(
                combinedAmount: combinedPersAmount,
                rothPercentage: rothPercentage,
                rothPercentageCompany: rothPercentage,
                combinedAmountCompany: combinedAllAmount
            )
    
            monthlyContributions.append(contribution)
        }
    
        return monthlyContributions
    }
    
    var monthlyPay            : Decimal = 12758.0
    var personal401kLimit     : Decimal = 22500.0
    var personal401kPercentage: Decimal = 100.0
    var companyContributionPercentage: Decimal = 16.0
    var roth401kPercentage    : Decimal = 10.0
    var combinedContributionLimit : Decimal = 66000.0
    
    let monthlyContributions = computeMonthlyContributions(
        monthlyPay            : monthlyPay,
        contributionPercentage: personal401kPercentage,
        companyContributionPercentage : companyContributionPercentage,
        rothPercentage        : roth401kPercentage,
        contributionLimit     : personal401kLimit,
        combinedContributionLimit : combinedContributionLimit
    )
    
    for month in ContributionMonth.allCases
    {
        let curContribution = monthlyContributions[month]
        print("Contribution for \(month)")
        print("    Traditional : $\(curContribution.traditional)")
        print("    Roth        : $\(curContribution.roth)")
        print("    Year-To-Date: $\(curContribution.yearToDate)")
        print("    Company Trad: $\(curContribution.traditionalCompany)")
        print("    Company Roth: $\(curContribution.rothCompany)")
        print("    Combined Year-To-Date: $\(curContribution.yearToDateAll)")
        print()
    }

The output:

    Contribution for January
        Traditional : $11482.2
        Roth        : $1275.8
        Year-To-Date: $12758
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $14799.28
    
    Contribution for February
        Traditional : $8767.8
        Roth        : $974.2
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $26582.56
    
    Contribution for March
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $28623.84
    
    Contribution for April
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $30665.12
    
    Contribution for May
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $32706.4
    
    Contribution for June
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $34747.68
    
    Contribution for July
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $36788.96
    
    Contribution for August
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $38830.24
    
    Contribution for September
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $40871.52
    
    Contribution for October
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $42912.8
    
    Contribution for November
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $44954.08
    
    Contribution for December
        Traditional : $0
        Roth        : $0
        Year-To-Date: $22500
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $46995.36

Now on to improve this code to include 401a contributions...

This is a snipped image of a working excel product that contains the functionality I wish to translate into Playground:

spreadsheet snip

I have updated the above "foundation" code in an attempt to incorporate 401a contributions. You'll see that 401a contributions in January and February look correct (matches the excel spreadsheet) but March and on are not correct (do not match the excel spreadsheet):

    import Foundation
    
    let pennyRoundingBehavior = NSDecimalNumberHandler(
        roundingMode: .bankers,
        scale: 2,
        raiseOnExactness: false,
        raiseOnOverflow: true,
        raiseOnUnderflow: true,
        raiseOnDivideByZero: true
    )
    
    func roundToNearestPenny(percentage: Decimal, of dollarAmount: Decimal) -> Decimal
    {
        assert((0...100).contains(percentage))
        
        let x = ((dollarAmount * percentage / 100) as NSDecimalNumber)
        return x.rounding(accordingToBehavior: pennyRoundingBehavior) as Decimal
    }
    
    enum ContributionMonth: Int, CaseIterable
    {
        case January, February, March, April, May, June,
             July, August, September, October, November, December
        
        static var count:Int  { allCases.count }
    }
    
    struct MonthlyContributions
    {
        var roth                : Decimal = 0
        var traditional         : Decimal = 0
        var pers401a            : Decimal = 0
        var rothCompany         : Decimal = 0
        var traditionalCompany  : Decimal = 0
        var yearToDate          : Decimal = 0
        var yearToDateAll       : Decimal = 0
    }
    
    extension MonthlyContributions
    {
        var total: Decimal { roth + traditional }
        var companyTotal: Decimal { rothCompany + traditionalCompany }
        
        func nextMonth(
            combinedAmount: Decimal,
            rothPercentage: Decimal,
            rothPercentageCompany: Decimal,
            combinedAmountCompany: Decimal
            ) -> Self
        {
            let roth = roundToNearestPenny(
                percentage: rothPercentage,
                of: combinedAmount
            )
            let rothCompany = roundToNearestPenny(
                percentage: rothPercentageCompany,
                of: combinedAmountCompany
            )
            return Self(
                roth: roth,
                traditional: combinedAmount - roth,
                pers401a: combinedAmount < total ? total - combinedAmount : 0,
                rothCompany: rothCompany,
                traditionalCompany: combinedAmountCompany - rothCompany,
                yearToDate: combinedAmount + pers401a + yearToDate,
                yearToDateAll: combinedAmount + combinedAmountCompany + yearToDateAll
            )
        }
    }
    
    extension Array where Element == MonthlyContributions {
        subscript (index: ContributionMonth) -> Element { self[index.rawValue] }
    }
    
    func calculatePersCombinedContribution(
        monthlyContribution        monthly      : Decimal,
        annualContributionLimit    limit        : Decimal,
        contributionsSoFarThisYear cummulative  : Decimal) -> Decimal
    {
        min(max(0, limit - cummulative), monthly)
    }
    func calculate401aContribution(
        monthlyContribution         monthly     : Decimal,
        combinedContributionLimit    limit        : Decimal,
        contributionsSoFarThisYear cummulative  : Decimal) -> Decimal
    {
        min(max(0, limit - cummulative), monthly)
    }
    func calculateAllCombinedContribution(
        monthlyContribution             monthly : Decimal,
        annualCombinedContributionLimit limit   : Decimal,
        allContributionsSoFarThisYear   cummulative: Decimal) -> Decimal
    {
        min(max(0, limit - cummulative), monthly)
    }
    
    func computeMonthlyContributions(
        monthlyPay            : Decimal,
        contributionPercentage: Decimal,
        companyContributionPercentage: Decimal,
        rothPercentage        : Decimal,
        contributionLimit     : Decimal,
        combinedContributionLimit: Decimal) -> [MonthlyContributions]
    {
        assert((0...100).contains(contributionPercentage))
        assert((0...100).contains(rothPercentage))
        assert(monthlyPay >= 0)
        assert(contributionLimit >= 0)
    
        var contribution = MonthlyContributions()
    
        var monthlyContributions: [MonthlyContributions] = []
        monthlyContributions.reserveCapacity(ContributionMonth.count)
    
        let monthlyContribution = roundToNearestPenny(
            percentage: contributionPercentage,
            of: monthlyPay
        )
        let monthlyCompanyContribution = roundToNearestPenny(
            percentage: companyContributionPercentage,
            of: monthlyPay
        )
        for _ in ContributionMonth.allCases
        {
            let combinedPersAmount = calculatePersCombinedContribution(
                monthlyContribution       : monthlyContribution,
                annualContributionLimit   : contributionLimit,
                contributionsSoFarThisYear: contribution.yearToDate
            )
            let pers401Amount = calculate401aContribution(
                monthlyContribution: monthlyContribution,
                combinedContributionLimit: combinedContributionLimit,
                contributionsSoFarThisYear: contribution.yearToDateAll
            )
            let combinedAllAmount = calculateAllCombinedContribution(
                monthlyContribution             : monthlyCompanyContribution,
                annualCombinedContributionLimit : combinedContributionLimit,
                allContributionsSoFarThisYear   : contribution.yearToDateAll
            )
            contribution = contribution.nextMonth(
                combinedAmount: combinedPersAmount,
                rothPercentage: rothPercentage,
                rothPercentageCompany: rothPercentage,
                combinedAmountCompany: combinedAllAmount
            )
    
            monthlyContributions.append(contribution)
        }
    
        return monthlyContributions
    }
    
    var monthlyPay            : Decimal = 12758.0
    var personal401kLimit     : Decimal = 22500.0
    var personal401kPercentage: Decimal = 100.0
    var companyContributionPercentage: Decimal = 16.0
    var roth401kPercentage    : Decimal = 10.0
    var combinedContributionLimit : Decimal = 66000.0
    
    let monthlyContributions = computeMonthlyContributions(
        monthlyPay            : monthlyPay,
        contributionPercentage: personal401kPercentage,
        companyContributionPercentage : companyContributionPercentage,
        rothPercentage        : roth401kPercentage,
        contributionLimit     : personal401kLimit,
        combinedContributionLimit : combinedContributionLimit
    )
    
    for month in ContributionMonth.allCases
    {
        let curContribution = monthlyContributions[month]
        print("Contribution for \(month)")
        print("    Traditional : $\(curContribution.traditional)")
        print("    Roth        : $\(curContribution.roth)")
        print("    Pers 401a   : $\(curContribution.pers401a)")
    //    print("    Pers Year-To-Date: $\(curContribution.yearToDate)")
        print("    Company Trad: $\(curContribution.traditionalCompany)")
        print("    Company Roth: $\(curContribution.rothCompany)")
        print("    Combined Year-To-Date: $\(curContribution.yearToDateAll)")
        print()
    }

The output:

    Contribution for January
        Traditional : $11482.2
        Roth        : $1275.8
        Pers 401a   : $0
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $14799.28
    
    Contribution for February
        Traditional : $8767.8
        Roth        : $974.2
        Pers 401a   : $3016
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $26582.56
    
    Contribution for March
        Traditional : $0
        Roth        : $0
        Pers 401a   : $9742
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $28623.84
    
    Contribution for April
        Traditional : $0
        Roth        : $0
        Pers 401a   : $0
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $30665.12
    
    Contribution for May
        Traditional : $0
        Roth        : $0
        Pers 401a   : $0
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $32706.4
    
    Contribution for June
        Traditional : $0
        Roth        : $0
        Pers 401a   : $0
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $34747.68
    
    Contribution for July
        Traditional : $0
        Roth        : $0
        Pers 401a   : $0
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $36788.96
    
    ...
    Contribution for December
        Traditional : $0
        Roth        : $0
        Pers 401a   : $0
        Company Trad: $1837.15
        Company Roth: $204.13
        Combined Year-To-Date: $46995.36

The goal is for the printouts for "Pers 401a" and "Combined Year-To-Date" monthly values to reflect the excel spreadsheet.

UPDATE

I created an additional year-to-date totals with the purpose to correctly limit the 401a, but then realized that the conditionals which define 401a only reference two types of year-to-date totals: 1) the sum of the previous month personal Roth and Traditional contributions & 2) the sum of the previous month personal, company and 401a contributions. Consequently, I believe I only need "yearToDate" and "yearToDateAll" variables so long as the "yearToDateAll" includes previous month 401a contributions.

I have outlined the conditionals which define each month's 401a contributions in comments below and applied the comments into code, but the results are messy and inaccurate.


import Foundation

let pennyRoundingBehavior = NSDecimalNumberHandler(
    roundingMode: .bankers,
    scale: 2,
    raiseOnExactness: false,
    raiseOnOverflow: true,
    raiseOnUnderflow: true,
    raiseOnDivideByZero: true
)

func roundToNearestPenny(percentage: Decimal, of dollarAmount: Decimal) -> Decimal
{
    assert((0...100).contains(percentage))
    
    let x = ((dollarAmount * percentage / 100) as NSDecimalNumber)
    return x.rounding(accordingToBehavior: pennyRoundingBehavior) as Decimal
}

enum ContributionMonth: Int, CaseIterable
{
    case January, February, March, April, May, June,
         July, August, September, October, November, December
    
    static var count:Int  { allCases.count }
}

struct MonthlyContributions
{
    var roth                : Decimal = 0
    var traditional         : Decimal = 0
    var rothCompany         : Decimal = 0
    var traditionalCompany  : Decimal = 0
    var yearToDate          : Decimal = 0
    var yearToDateAll       : Decimal = 0
    var pers401a            : Decimal = 0
}

extension MonthlyContributions
{
    var total: Decimal { roth + traditional }
    var companyTotal: Decimal { rothCompany + traditionalCompany }
    
    func nextMonth(
        combinedAmount: Decimal,
        rothPercentage: Decimal,
        rothPercentageCompany: Decimal,
        combinedAmountCompany: Decimal
        ) -> Self
    {
        let roth = roundToNearestPenny(
            percentage: rothPercentage,
            of: combinedAmount
        )
        let rothCompany = roundToNearestPenny(
            percentage: rothPercentageCompany,
            of: combinedAmountCompany
        )
        return Self(
            roth: roth,
            traditional: combinedAmount - roth,
            rothCompany: rothCompany,
            traditionalCompany: combinedAmountCompany - rothCompany,
            yearToDate: combinedAmount + yearToDate,
            yearToDateAll: combinedAmount + pers401a + combinedAmountCompany + yearToDateAll,
            /*
             pers401a conditionals:
             1. If the sum of the previous month's running total of 
                personal and company contributions + this month's 
                personal contributions is < personal401klimit 
                ($22,500), then pers401a = 0, otherwise...
             2. If the sum of the previous month's running total of 
                personal and company contributions >= the 
                combined401klimit ($66,500), then pers401a = 0, 
                otherwise...
             3. If the sum of the previous month's running total of 
                personal and company contributions + this month's 
                personal roth and traditional contributions < the 
                combined401klimit ($66,500), then...
                    a. if ((this month's personal roth + 
                       traditional contributions) - (sum of all                
                       previous month's peresonal roth & traditional
                       contributions)) = 0, then pers401a = 0, otherwise...
                    b. pers401a = (this month's personal roth + traditional
                       contributions) - (sum of all previous month's 
                       peresonal roth & traditional contributions), otherwise...
             4. pers401a = combined401klimit - the sum of the previous 
                month's running total of personal and company  contributions.
            */
            pers401a:
                yearToDateAll + total < personal401kLimit ? 0 :
                yearToDateAll >= combinedContributionLimit ? 0 :
                yearToDateAll + total < combinedContributionLimit ?
                    (total - yearToDate) == 0 ? 0 :
                    (total - yearToDate) :
                combinedContributionLimit - yearToDateAll
        )
    }
}

extension Array where Element == MonthlyContributions {
    subscript (index: ContributionMonth) -> Element { self[index.rawValue] }
}

func calculateAllCombinedContribution(
    monthlyContribution          monthly    : Decimal,
    annualContributionLimit      limit      : Decimal,
    contributionsSoFarThisYear   cummulative: Decimal) -> Decimal
{
    min(max(0, limit - cummulative), monthly)
}

func computeMonthlyContributions(
    monthlyPay                      : Decimal,
    contributionPercentage          : Decimal,
    companyContributionPercentage   : Decimal,
    rothPercentage                  : Decimal,
    contributionLimit               : Decimal,
    combinedContributionLimit       : Decimal) -> [MonthlyContributions]
{
    assert((0...100).contains(contributionPercentage))
    assert((0...100).contains(rothPercentage))
    assert(monthlyPay >= 0)
    assert(contributionLimit >= 0)

    var contribution = MonthlyContributions()

    var monthlyContributions: [MonthlyContributions] = []
    monthlyContributions.reserveCapacity(ContributionMonth.count)

    let monthlyContribution = roundToNearestPenny(
        percentage: contributionPercentage,
        of: monthlyPay
    )
    let monthlyCompanyContribution = roundToNearestPenny(
        percentage: companyContributionPercentage,
        of: monthlyPay
    )
    for _ in ContributionMonth.allCases
    {
        let combinedPersAmount = calculateAllCombinedContribution(
            monthlyContribution       : monthlyContribution,
            annualContributionLimit   : contributionLimit,
            contributionsSoFarThisYear: contribution.yearToDate
        )
        let pers401Amount = calculateAllCombinedContribution(
            monthlyContribution         : monthlyContribution,
            annualContributionLimit     : combinedContributionLimit,
            //incorporate new yearToDate401a
            contributionsSoFarThisYear  : contribution.yearToDateAll
        )
        let combinedAllAmount = calculateAllCombinedContribution(
            monthlyContribution         : monthlyCompanyContribution,
            annualContributionLimit     : combinedContributionLimit,
            contributionsSoFarThisYear  : contribution.yearToDateAll
        )
        contribution = contribution.nextMonth(
            combinedAmount: combinedPersAmount,
            rothPercentage: rothPercentage,
            rothPercentageCompany: rothPercentage,
            combinedAmountCompany: combinedAllAmount
        )

        monthlyContributions.append(contribution)
    }

    return monthlyContributions
}

var monthlyPay            : Decimal = 12758.0
var personal401kLimit     : Decimal = 22500.0
var personal401kPercentage: Decimal = 100.0
var companyContributionPercentage: Decimal = 16.0
var roth401kPercentage    : Decimal = 10.0
var combinedContributionLimit : Decimal = 66000.0

let monthlyContributions = computeMonthlyContributions(
    monthlyPay            : monthlyPay,
    contributionPercentage: personal401kPercentage,
    companyContributionPercentage : companyContributionPercentage,
    rothPercentage        : roth401kPercentage,
    contributionLimit     : personal401kLimit,
    combinedContributionLimit : combinedContributionLimit
)

for month in ContributionMonth.allCases
{
    let curContribution = monthlyContributions[month]
    print("Contribution for \(month)")
    print("    Traditional : $\(curContribution.traditional)")
    print("    Roth        : $\(curContribution.roth)")
    print("    Pers 401a   : $\(curContribution.pers401a)")
//    print("    Pers Year-To-Date: $\(curContribution.yearToDate)")
    print("    Company Trad: $\(curContribution.traditionalCompany)")
    print("    Company Roth: $\(curContribution.rothCompany)")
    print("    Combined Year-To-Date: $\(curContribution.yearToDateAll)")
    print()
}

The output:

Contribution for January
    Traditional : $11482.2
    Roth        : $1275.8
    Pers 401a   : $0
    Company Trad: $1837.15
    Company Roth: $204.13
    Combined Year-To-Date: $14799.28

Contribution for February
    Traditional : $8767.8
    Roth        : $974.2
    Pers 401a   : $0
    Company Trad: $1837.15
    Company Roth: $204.13
    Combined Year-To-Date: $26582.56

Contribution for March
    Traditional : $0
    Roth        : $0
    Pers 401a   : $-12758
    Company Trad: $1837.15
    Company Roth: $204.13
    Combined Year-To-Date: $28623.84

Contribution for April
    Traditional : $0
    Roth        : $0
    Pers 401a   : $-22500
    Company Trad: $1837.15
    Company Roth: $204.13
    Combined Year-To-Date: $17907.12

Contribution for May
    Traditional : $0
    Roth        : $0
    Pers 401a   : $0
    Company Trad: $1837.15
    Company Roth: $204.13
    Combined Year-To-Date: $-2551.6


...
Solution

  • I used your last version of the code as my starting point.

    I think that you are trying to do too much in the nextMonth() method. You compute the 401k contribution amount before calling it. I would suggest computing the 401a contribution the same way. So nextMonth() can be simplified.

    extension MonthlyContributions
{
    var total: Decimal { roth + traditional }
    var companyTotal: Decimal { rothCompany + traditionalCompany }
    
    func nextMonth(
        combinedAmount: Decimal,
        rothPercentage: Decimal,
        rothPercentageCompany: Decimal,
        combinedAmountCompany: Decimal,
        pers401a: Decimal) -> Self
    {
        let roth = roundToNearestPenny(
            percentage: rothPercentage,
            of: combinedAmount
        )
        let rothCompany = roundToNearestPenny(
            percentage: rothPercentageCompany,
            of: combinedAmountCompany
        )
        return Self(
            roth: roth,
            traditional: combinedAmount - roth,
            rothCompany: rothCompany,
            traditionalCompany: combinedAmountCompany - rothCompany,
            yearToDate: combinedAmount + yearToDate,
            yearToDateAll: combinedAmount + pers401a + combinedAmountCompany + yearToDateAll,
            pers401a: pers401a
        )
    }
}

    In computeMonthlyContributions there are a couple of problems.

    The first is that you are computing the pre-limit company contribution outside the loop from constant pre-limit employee contributions. I think you were just following the same pattern for company contributions as for employee contributions, but that's not quite right. The amount the company contributes should be computed based on the employee's actual contribution after the employee limit has been applied, and that's only available in the loop body.

    Another issue is that when applying the annual limit to the 401a contribution, both personal and company contributions for the current month have to be included in the year-to-date total.

    I renamed the function to limit contributions to be a little clearer - or at least I hope it's clearer:

    func clampContribution(
    _                   current    : Decimal,
    whenCumulativeValue cummulative: Decimal,
    exceeds             limit      : Decimal) -> Decimal
{
    min(max(0, limit - cummulative), current)
}

    Note that I swapped the order of the last two parameters compared to the previous version. That was to make it read more like an English sentence. It does exactly the same thing as the previous version, so it's just a matter of naming. If you keep the previous one, be aware that the code below uses this new version, so you'll need to appropriately swap the last two parameters where I call clampContributions.

    Then made the changes I described to computeMonthlyContributions:

    func computeMonthlyContributions(
    monthlyPay                      : Decimal,
    contributionPercentage          : Decimal,
    companyContributionPercentage   : Decimal,
    rothPercentage                  : Decimal,
    contributionLimit               : Decimal,
    combinedContributionLimit       : Decimal) -> [MonthlyContributions]
{
    assert((0...100).contains(contributionPercentage))
    assert((0...100).contains(rothPercentage))
    assert(monthlyPay >= 0)
    assert(contributionLimit >= 0)

    var contribution = MonthlyContributions()

    var monthlyContributions: [MonthlyContributions] = []
    monthlyContributions.reserveCapacity(ContributionMonth.count)

    let monthlyContribution = roundToNearestPenny(
        percentage: contributionPercentage,
        of: monthlyPay
    )
    for _ in ContributionMonth.allCases
    {
        let combinedPersAmount = clampContribution(
            monthlyContribution,
            whenCumulativeValue: contribution.yearToDate,
            exceeds: contributionLimit
        )

        // ADDED THIS
        // Include personal contribution in overall YTD 401k total so it
        // will limit company contribution.
        var newYearToDateAll = contribution.yearToDateAll + combinedPersAmount

        // MOVED THIS FROM OUTSIDE THE LOOP TO HERE
        // Compute company contribution based on employee's actual contribution
        let monthlyCompanyContribution = roundToNearestPenny(
            percentage: companyContributionPercentage,
            of: combinedPersAmount // <-- CHANGED THIS
        )

        // Limit company contribution based on overall YTD 401k total,
        // which includes the employee's contribution for this month.
        let companyContribution = clampContribution(
            monthlyCompanyContribution,
            whenCumulativeValue: newYearToDateAll, // <-- CHANGED THIS
            exceeds: combinedContributionLimit
        )
        
        // ADDED THIS
        // Include company contribution in overall YTD 401k total so it 
        // will limit 401a contributions.
        newYearToDateAll += companyContribution
        
        // ADDED THIS
        // Compute amount of personal contribution available for 401a
        let preLimit401aAmount = monthlyContribution - combinedPersAmount;

        // Limit the 401a contribution based on overall YTD 401k total, 
        // which now includes this month's employee and company contributions.
        let pers401Amount = clampContribution(
            preLimit401aAmount,  // <-- CHANGED THIS
            whenCumulativeValue: newYearToDateAll, // <-- CHANGED THIS
            exceeds: combinedContributionLimit
        )

        contribution = contribution.nextMonth(
            combinedAmount: combinedPersAmount,
            rothPercentage: rothPercentage,
            rothPercentageCompany: rothPercentage,
            combinedAmountCompany: companyContribution,
            pers401a: pers401Amount
        )
        
        assert(contribution.yearToDate <= contributionLimit)
        assert(contribution.yearToDateAll <= combinedContributionLimit)

        monthlyContributions.append(contribution)
    }

    return monthlyContributions
}

    I also added sanity checks (assert calls) to verify that the limits aren't exceeded after the new MonthlyContributions instance is created.

    This is the output:

    Contribution for January
    Traditional : $11482.2
    Roth        : $1275.8
    Pers 401a   : $0
    Company Trad: $1837.15
    Company Roth: $204.13
    Combined Year-To-Date: $14799.28

Contribution for February
    Traditional : $8767.8
    Roth        : $974.2
    Pers 401a   : $3016
    Company Trad: $1402.85
    Company Roth: $155.87
    Combined Year-To-Date: $29116

Contribution for March
    Traditional : $0
    Roth        : $0
    Pers 401a   : $12758
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $41874

Contribution for April
    Traditional : $0
    Roth        : $0
    Pers 401a   : $12758
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $54632

Contribution for May
    Traditional : $0
    Roth        : $0
    Pers 401a   : $11368
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $66000

Contribution for June
    Traditional : $0
    Roth        : $0
    Pers 401a   : $0
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $66000

Contribution for July
    Traditional : $0
    Roth        : $0
    Pers 401a   : $0
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $66000

Contribution for August
    Traditional : $0
    Roth        : $0
    Pers 401a   : $0
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $66000

Contribution for September
    Traditional : $0
    Roth        : $0
    Pers 401a   : $0
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $66000

Contribution for October
    Traditional : $0
    Roth        : $0
    Pers 401a   : $0
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $66000

Contribution for November
    Traditional : $0
    Roth        : $0
    Pers 401a   : $0
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $66000

Contribution for December
    Traditional : $0
    Roth        : $0
    Pers 401a   : $0
    Company Trad: $0
    Company Roth: $0
    Combined Year-To-Date: $66000

    Refactored Version

    With the code working, it might be time to think how it would be used in an actual app. There are several things that it just doesn't handle. Employees can periodically change their contribution percentages. Their pay can change, hopefully for the better!

    Originally, I said that you were doing too much in nextMonth. It seemed to be causing confusion about how to do the calculations, especially since part of the calculations were done before calling nextMonth and part were done inside of nextMonth. I felt it was important to keep nextMonth simple, so all it did was split 401k contributions between Roth and traditional 401k.

    But now that it's working, given that it will need to be more flexible to handle the cases I mentioned above, it would be better to have all of the contribution code in one place. That simplifies the code that calls it, as well as encodes the knowledge of how to do the computation in one place. These changes in code structure are a normal part of programming. When you're just trying to get something work in the basic case, you might need it structured one way to help you think through it, but then once you have that working, and start thinking of its broader use, you might realize you need it structured differently.

    The problem with putting all that contribution code in nextMonth is that you need to pass in all the percentages and annual limits, and long parameter lists kind of suck. One solution is to create a "parameter pack" that contains all the parameter data. It's just a struct that you pass as one parameter, instead passing a lot of individual parameters.

    struct ContributionSplitInfo
{
    // These are static because they are shared among all employees.
    static let personal401kLimit        : Decimal = 22500.0
    static let combined401kAnnualLimit  : Decimal = 66000.0
    static let companyMatchingPercentage: Decimal = 16.0
    
    /*
     These are provided as convenience, and because if for some reason they
     need to be individualized for each employee, that can be done without
     affecting the code that uses it.  For example, tax law could change limits,
     but existing employees might be grandfathered with the old limits.
     Similarly the company might change its matching rate for new hires, but
     not for new hires, or might have increasing matching with years of service.
     */
    var personal401kAnnualLimit  : Decimal { Self.personal401kLimit }
    var combined401kAnnualLimit  : Decimal { Self.combined401kAnnualLimit }
    var companyMatchingPercentage: Decimal { Self.companyMatchingPercentage }
    
    // These are stored instance properties because each employee can set them
    // differently
    var personal401kPercentage   : Decimal
    var rothPercentage           : Decimal
}

    Note that it does not include the monthly pay. We'll only be using one instance for now, but you could create one for each employee, or even have multiple ones for each employee as they make changes to their percentages throughout the year.

    nextMonth needs changing to use it. I've added in all of the code to compute the current month's contributions, and extracted some of the individual computations into private methods, in hopes of making it more self-documenting. I did include one doc comment for one of the private methods, because the way I implemented it, it updates a year-to-date inout parameter, and I wanted to be clear what is expected on entry and what it will be on return.

    The following code goes in the MonthlyContributions extension.

        func nextMonth(pay: Decimal, using splitInfo: ContributionSplitInfo) -> Self
    {
        assert((0...100).contains(splitInfo.personal401kPercentage))
        assert((0...100).contains(splitInfo.companyMatchingPercentage))
        assert((0...100).contains(splitInfo.rothPercentage))
        assert(splitInfo.personal401kAnnualLimit >= 0)
        assert(splitInfo.combined401kAnnualLimit >= 0)
        assert(pay >= 0)


        let (personal401kContribution, prelimit401aContribution) =
            self.personal401kContribution(fromPay: pay, using: splitInfo)
        
        var combinedYearToDateContributions = yearToDateAll

        let company401kContribution = self.company401kContribution(
            fromPersonalContribution: personal401kContribution,
            andCombinedYearToDateContributions: &combinedYearToDateContributions,
            using: splitInfo
        )

        let personal401aContribution = clampContribution(
            prelimit401aContribution,
            whenCumulativeValue: combinedYearToDateContributions,
            exceeds: splitInfo.combined401kAnnualLimit
        )
        
        combinedYearToDateContributions += personal401aContribution
        
        let (personalTraditional, personalRoth) = split401kContribution(
            contribution: personal401kContribution,
            rothPercentage: splitInfo.rothPercentage
        )
        
        let (companyTraditional, companyRoth) = split401kContribution(
            contribution: company401kContribution,
            rothPercentage: splitInfo.rothPercentage
        )
        
        let contributions = Self(
            roth              : personalRoth,
            traditional       : personalTraditional,
            rothCompany       : companyRoth,
            traditionalCompany: companyTraditional,
            yearToDate        : yearToDate + personal401kContribution,
            yearToDateAll     : combinedYearToDateContributions,
            pers401a          : personal401aContribution
        )
        
        assert(contributions.yearToDate <= splitInfo.personal401kAnnualLimit)
        assert(contributions.yearToDateAll <= splitInfo.combined401kAnnualLimit)

        return contributions
    }
    
    private func split401kContribution(
        contribution: Decimal,
        rothPercentage: Decimal) -> (traditional: Decimal, roth: Decimal)
    {
        let roth = roundToNearestPenny(
            percentage: rothPercentage,
            of: contribution
        )
        
        return (contribution - roth, roth)
    }
    
    private func personal401kContribution(
        fromPay pay: Decimal,
        using splitInfo: ContributionSplitInfo)
        -> (contribution: Decimal, unused: Decimal)
    {
        let prelimitContribution = roundToNearestPenny(
            percentage: splitInfo.personal401kPercentage,
            of: pay
        )
        let actualContribution = clampContribution(
            prelimitContribution,
            whenCumulativeValue: yearToDate,
            exceeds: splitInfo.personal401kAnnualLimit
        )
        
        return (actualContribution, prelimitContribution - actualContribution)
    }
    
    /**
     Compute the company's matching 401k contribution for the current month.
          
     - Parameters:
        - personalContribution: The current month's actual personal 401k
            contribution, that is *after* the personal annual limit has been
            applied.
        - yearToDate: The year-to-date sum of all employee and company 401k
            contributions
            - On entry: The sum should  *not* include any information for the
                current month.
            - On return: `yearToDate` will be updated to include the current
                months personal and company contributions.
        - splitInfo: `ContributionSplitInfo` instance specifying the company
            matching percentage and annual limits.
     
     - Returns: The company's matching 401k contribution for the current month.
     */
    private func company401kContribution(
        fromPersonalContribution personalContribution: Decimal,
        andCombinedYearToDateContributions yearToDate: inout Decimal,
        using splitInfo: ContributionSplitInfo) -> Decimal
    {
        yearToDate += personalContribution
        
        let companyCombinedContribution = roundToNearestPenny(
            percentage: splitInfo.companyMatchingPercentage,
            of: personalContribution
        )

        let companyContribution = clampContribution(
            companyCombinedContribution,
            whenCumulativeValue: yearToDate,
            exceeds: splitInfo.combined401kAnnualLimit
        )
        
        yearToDate += companyContribution
        
        return companyContribution
    }

    Then computeMonthlyContributions becomes super simple, which was a large part of the reason for this refactoring in the first place. Now whatever code that is dealing with months or some data structure for holding the MonthlyContributions instances doesn't need to know anything at all about the details of computing the contributions.

    func computeMonthlyContributions(
    monthlyPay          : Decimal,
    using      splitInfo: ContributionSplitInfo) -> [MonthlyContributions]
{
    var contribution = MonthlyContributions()

    var monthlyContributions: [MonthlyContributions] = []
    monthlyContributions.reserveCapacity(ContributionMonth.count)

    for _ in ContributionMonth.allCases
    {
        contribution = contribution.nextMonth(pay: monthlyPay, using: splitInfo)
        monthlyContributions.append(contribution)
    }

    return monthlyContributions
}

    All the globals you had defining the percentages and limits are changed to this, along with the code to call computeMonthlyContributions:

    var monthlyPay: Decimal = 12758.0
let splitInfo =
    ContributionSplitInfo(personal401kPercentage: 100.0, rothPercentage: 16.0)

let monthlyContributions =
    computeMonthlyContributions(monthlyPay: monthlyPay, using: splitInfo)

    The print loop code remains unchanged.

    Now if you want, you can experiment with changing the pay in the loop, or with changing the percentages from month to month in a way that can happen in the real world.