feat(phase-b): persist context bigrams in NSUserDefaults (capped); add persistence test; minor wiring

Author: arnabwahid

ContextRanking.m

@@ -4,17 +4,84 @@
 //
 
 #import "ContextRanking.h"
+#import <Foundation/Foundation.h>
 
 static NSMutableArray<NSString *> *s_recentHistory;
 static NSMutableDictionary<NSString *, NSMutableDictionary<NSString *, NSNumber *> *> *s_bigrams; // prev -> (next -> count)
 static const NSUInteger kHistoryMax = 8; // small cap to keep things light
+static NSString * const kBigramsDefaultsKey = @"ContextRankingBigrams";
+static const NSUInteger kPrevMax = 64;      // cap number of prev tokens to persist
+static const NSUInteger kNextPerPrevMax = 12; // cap next variants per prev
+
+static void ensureStoresLoaded(void) {
+    if (!s_bigrams) {
+        NSDictionary *saved = [[NSUserDefaults standardUserDefaults] dictionaryForKey:kBigramsDefaultsKey];
+        s_bigrams = [[NSMutableDictionary alloc] init];
+        if ([saved isKindOfClass:[NSDictionary class]]) {
+            [saved enumerateKeysAndObjectsUsingBlock:^(NSString *prev, NSDictionary *nexts, BOOL *stop) {
+                if (![prev isKindOfClass:[NSString class]] || ![nexts isKindOfClass:[NSDictionary class]]) return;
+                NSMutableDictionary<NSString *, NSNumber *> *mutableNexts = [[NSMutableDictionary alloc] init];
+                [nexts enumerateKeysAndObjectsUsingBlock:^(NSString *nxt, NSNumber *cnt, BOOL *stop2) {
+                    if ([nxt isKindOfClass:[NSString class]] && [cnt isKindOfClass:[NSNumber class]]) {
+                        mutableNexts[nxt] = cnt;
+                    }
+                }];
+                s_bigrams[prev] = mutableNexts;
+            }];
+        }
+    }
+}
+
+static void persistBigrams(void) {
+    if (!s_bigrams) return;
+    // Enforce caps: limit number of prev keys and next variants per prev (by lowest count pruning)
+    if (s_bigrams.count > kPrevMax) {
+        NSArray *keys = [s_bigrams allKeys];
+        NSMutableArray *scored = [NSMutableArray arrayWithCapacity:keys.count];
+        for (NSString *k in keys) {
+            NSDictionary *m = s_bigrams[k];
+            NSInteger sum = 0; for (NSNumber *v in [m allValues]) sum += v.integerValue;
+            [scored addObject:@{ @"k": k, @"s": @(sum) }];
+        }
+        [scored sortUsingComparator:^NSComparisonResult(NSDictionary *a, NSDictionary *b){
+            return [a[@"s"] integerValue] < [b[@"s"] integerValue] ? NSOrderedAscending : NSOrderedDescending;
+        }];
+        for (NSUInteger i = kPrevMax; i < scored.count; i++) {
+            NSString *drop = scored[i][@"k"];
+            [s_bigrams removeObjectForKey:drop];
+        }
+    }
+    for (NSString *prev in [s_bigrams allKeys]) {
+        NSMutableDictionary<NSString *, NSNumber *> *nextMap = s_bigrams[prev];
+        if (nextMap.count > kNextPerPrevMax) {
+            NSArray *pairs = [nextMap allKeys];
+            NSArray *sorted = [pairs sortedArrayUsingComparator:^NSComparisonResult(NSString *a, NSString *b){
+                NSInteger ca = [nextMap[a] integerValue];
+                NSInteger cb = [nextMap[b] integerValue];
+                if (ca == cb) return NSOrderedSame;
+                return (ca > cb) ? NSOrderedAscending : NSOrderedDescending;
+            }];
+            NSSet *keep = [NSSet setWithArray:[sorted subarrayWithRange:NSMakeRange(0, kNextPerPrevMax)]];
+            for (NSString *key in [nextMap allKeys]) {
+                if (![keep containsObject:key]) [nextMap removeObjectForKey:key];
+            }
+        }
+    }
+    NSMutableDictionary *toSave = [NSMutableDictionary dictionaryWithCapacity:s_bigrams.count];
+    [s_bigrams enumerateKeysAndObjectsUsingBlock:^(NSString *prev, NSDictionary *nexts, BOOL *stop){
+        toSave[prev] = [nexts copy];
+    }];
+    [[NSUserDefaults standardUserDefaults] setObject:toSave forKey:kBigramsDefaultsKey];
+    [[NSUserDefaults standardUserDefaults] synchronize];
+}
 
 @implementation ContextRanking
 
 + (NSArray<NSString *> *)rankCandidates:(NSArray<NSString *> *)candidates
                            withHistory:(NSArray<NSString *> * _Nullable)history
 {
     if (!candidates || candidates.count <= 1) return candidates ?: @[];
+    ensureStoresLoaded();
     // Minimal heuristic: boost candidates observed to follow the last committed token.
     NSString *last = (history.count > 0 ? history[0] : nil);
     if (!last || !s_bigrams) {
@@ -43,11 +110,11 @@ @implementation ContextRanking
 + (void)recordCommittedToken:(NSString *)token {
     if (token.length == 0) return;
     @synchronized(self) {
+        ensureStoresLoaded();
         if (!s_recentHistory) s_recentHistory = [[NSMutableArray alloc] init];
         // Update bigram counts using previous token (if any)
         NSString *prev = (s_recentHistory.count > 0 ? s_recentHistory[0] : nil);
         if (prev && token) {
-            if (!s_bigrams) s_bigrams = [[NSMutableDictionary alloc] init];
             NSMutableDictionary<NSString *, NSNumber *> *nextMap = s_bigrams[prev];
             if (!nextMap) {
                 nextMap = [[NSMutableDictionary alloc] init];
@@ -62,6 +129,7 @@ + (void)recordCommittedToken:(NSString *)token {
         [s_recentHistory insertObject:token atIndex:0];
         while (s_recentHistory.count > kHistoryMax) [s_recentHistory removeLastObject];
     }
+    persistBigrams();
 }
 
 + (NSArray<NSString *> *)recentHistory:(NSUInteger)limit {

Tests/ContextRankingTests.m

@@ -39,4 +39,20 @@ - (void)testBigramBoostMovesSeenPairUp {
     XCTAssertEqualObjects(out[0], curr, @"Observed bigram should be boosted to front");
 }
 
+- (void)testPersistenceWritesBigramsToUserDefaults {
+    // Clear any prior data
+    [[NSUserDefaults standardUserDefaults] removeObjectForKey:@"ContextRankingBigrams"];
+    [[NSUserDefaults standardUserDefaults] synchronize];
+    NSString *prev = @"alpha";
+    NSString *next = @"beta";
+    [ContextRanking recordCommittedToken:prev];
+    [ContextRanking recordCommittedToken:next];
+    NSDictionary *saved = [[NSUserDefaults standardUserDefaults] dictionaryForKey:@"ContextRankingBigrams"];
+    XCTAssertNotNil(saved);
+    NSDictionary *nexts = saved[prev];
+    XCTAssertTrue([nexts isKindOfClass:[NSDictionary class]]);
+    NSNumber *cnt = nexts[next];
+    XCTAssertTrue(cnt != nil && cnt.integerValue >= 1);
+}
+
 @end