dead_code_detector.py

"""
Dead/Rarely-Used Code Detector using Deep Learning
Binary Classification: 0 = frequently used, 1 = rarely/never used
"""

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score, roc_curve
import matplotlib.pyplot as plt
import seaborn as sns
import re
from collections import defaultdict
import networkx as nx


class CallGraphParser:
    """Parse the callgraph.txt file and extract features"""
    
    def __init__(self, filepath):
        self.filepath = filepath
        self.graph = nx.DiGraph()
        self.function_data = {}
        
    def parse(self):
        """Parse the callgraph file"""
        print("Parsing callgraph...")
        
        with open(self.filepath, 'r') as f:
            lines = f.readlines()
        
        current_function = None
        current_uses = 0
        calls = []
        
        for line in lines:
            line = line.strip()
            
            # Match call graph node lines
            if line.startswith("Call graph node"):
                # Save previous function if exists
                if current_function is not None:
                    self.function_data[current_function] = {
                        'uses': current_uses,
                        'calls': calls,
                        'num_calls': len(calls)
                    }
                
                # Extract function name and uses
                match = re.search(r"for function: '([^']+)'.*#uses=(\d+)", line)
                if match:
                    current_function = match.group(1)
                    current_uses = int(match.group(2))
                else:
                    # Handle null function or other formats
                    match = re.search(r"<<null function>>.*#uses=(\d+)", line)
                    if match:
                        current_function = "<<null>>"
                        current_uses = int(match.group(1))
                    else:
                        current_function = None
                        current_uses = 0
                
                calls = []
            
            # Match function calls
            elif line.startswith("CS<") and current_function:
                match = re.search(r"calls function '([^']+)'", line)
                if match:
                    called_func = match.group(1)
                    calls.append(called_func)
                    self.graph.add_edge(current_function, called_func)
        
        # Save last function
        if current_function is not None:
            self.function_data[current_function] = {
                'uses': current_uses,
                'calls': calls,
                'num_calls': len(calls)
            }
        
        print(f"Parsed {len(self.function_data)} functions")
        return self.function_data
    
    def extract_features(self, threshold=2):
        """
        Extract features for each function
        threshold: functions with uses <= threshold are labeled as rarely used (1)
        """
        print("\nExtracting features...")
        
        features_list = []
        
        for func_name, data in self.function_data.items():
            # Basic features
            uses = data['uses']
            num_calls = data['num_calls']
            
            # Graph-based features
            try:
                # How many functions call this function = in-degree
                in_degree = self.graph.in_degree(func_name) if func_name in self.graph else 0
                
                # How many functions this calls = out-degree
                out_degree = self.graph.out_degree(func_name) if func_name in self.graph else 0
                
                # PageRank centrality
                pagerank = 0
                
                # Betweenness centrality (computationally expensive, optional)
                betweenness = 0
                
            except:
                in_degree = 0
                out_degree = num_calls
                pagerank = 0
                betweenness = 0
            
            # Function name features
            func_length = len(func_name)
            has_number = int(any(c.isdigit() for c in func_name))
            has_underscore = int('_' in func_name)
            is_internal = int(func_name.startswith('_') or '<<' in func_name)
            has_sqlite_prefix = int(func_name.startswith('sqlite3'))
            
            # Label: 1 if rarely used (uses <= threshold), 0 otherwise
            label = 1 if uses <= threshold else 0
            
            features = {
                'function_name': func_name,
                'uses': uses,
                'num_calls': num_calls,
                'in_degree': in_degree,
                'out_degree': out_degree,
                'pagerank': pagerank,
                'betweenness': betweenness,
                'func_length': func_length,
                'has_number': has_number,
                'has_underscore': has_underscore,
                'is_internal': is_internal,
                'has_sqlite_prefix': has_sqlite_prefix,
                'label': label
            }
            
            features_list.append(features)
        
        df = pd.DataFrame(features_list)
        
        # Calculate PageRank if graph is not empty
        if len(self.graph.nodes()) > 0:
            print("Computing PageRank...")
            try:
                pagerank_dict = nx.pagerank(self.graph, max_iter=100)
                df['pagerank'] = df['function_name'].map(lambda x: pagerank_dict.get(x, 0))
            except:
                print("PageRank computation failed, using zeros")
        
        print(f"\nDataset statistics:")
        print(f"Total functions: {len(df)}")
        print(f"Rarely used (label=1): {(df['label']==1).sum()} ({(df['label']==1).sum()/len(df)*100:.1f}%)")
        print(f"Frequently used (label=0): {(df['label']==0).sum()} ({(df['label']==0).sum()/len(df)*100:.1f}%)")
        
        return df


class CodeDataset(Dataset):
    """PyTorch Dataset for code classification"""
    
    def __init__(self, features, labels):
        self.features = torch.FloatTensor(features)
        self.labels = torch.LongTensor(labels)
    
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]


class DeadCodeClassifier(nn.Module):
    """Deep Learning model for dead code detection"""
    
    def __init__(self, input_dim, hidden_dims=[64, 32, 16], dropout=0.3):
        super(DeadCodeClassifier, self).__init__()
        
        layers = []
        prev_dim = input_dim
        
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout))
            prev_dim = hidden_dim
        
        layers.append(nn.Linear(prev_dim, 2))  # Binary classification
        
        self.network = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.network(x)


class Trainer:
    """Training pipeline"""
    
    def __init__(self, model, device='cpu'):
        self.model = model.to(device)
        self.device = device
        self.history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
    
    def train_epoch(self, dataloader, optimizer, criterion):
        self.model.train()
        total_loss = 0
        correct = 0
        total = 0
        
        for features, labels in dataloader:
            features, labels = features.to(self.device), labels.to(self.device)
            
            optimizer.zero_grad()
            outputs = self.model(features)
            loss = criterion(outputs, labels)
            
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
        
        return total_loss / len(dataloader), 100. * correct / total
    
    def validate(self, dataloader, criterion):
        self.model.eval()
        total_loss = 0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for features, labels in dataloader:
                features, labels = features.to(self.device), labels.to(self.device)
                
                outputs = self.model(features)
                loss = criterion(outputs, labels)
                
                total_loss += loss.item()
                _, predicted = outputs.max(1)
                total += labels.size(0)
                correct += predicted.eq(labels).sum().item()
        
        return total_loss / len(dataloader), 100. * correct / total
    
    def fit(self, train_loader, val_loader, epochs=100, lr=0.001, weight_decay=1e-5):
        criterion = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(self.model.parameters(), lr=lr, weight_decay=weight_decay)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', 
                                                                factor=0.5, patience=10)
        
        best_val_acc = 0
        patience_counter = 0
        patience = 20
        
        print("\nStarting training...")
        for epoch in range(epochs):
            train_loss, train_acc = self.train_epoch(train_loader, optimizer, criterion)
            val_loss, val_acc = self.validate(val_loader, criterion)
            
            self.history['train_loss'].append(train_loss)
            self.history['train_acc'].append(train_acc)
            self.history['val_loss'].append(val_loss)
            self.history['val_acc'].append(val_acc)
            
            scheduler.step(val_loss)
            
            if (epoch + 1) % 10 == 0:
                print(f'Epoch [{epoch+1}/{epochs}] '
                      f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}% | '
                      f'Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%')
            
            # Early stopping
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                patience_counter = 0
                torch.save(self.model.state_dict(), 'best_model.pth')
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    print(f'\nEarly stopping at epoch {epoch+1}')
                    break
        
        print(f'\nBest validation accuracy: {best_val_acc:.2f}%')
        self.model.load_state_dict(torch.load('best_model.pth'))
    
    def evaluate(self, dataloader):
        self.model.eval()
        all_preds = []
        all_labels = []
        all_probs = []
        
        with torch.no_grad():
            for features, labels in dataloader:
                features = features.to(self.device)
                outputs = self.model(features)
                probs = F.softmax(outputs, dim=1)
                _, predicted = outputs.max(1)
                
                all_preds.extend(predicted.cpu().numpy())
                all_labels.extend(labels.numpy())
                all_probs.extend(probs[:, 1].cpu().numpy())
        
        return np.array(all_labels), np.array(all_preds), np.array(all_probs)


def plot_results(history, y_true, y_pred, y_probs, class_names=['Frequently Used', 'Rarely Used']):
    """Plot training history and evaluation metrics"""
    
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    
    # Plot 1: Training history
    ax = axes[0, 0]
    ax.plot(history['train_loss'], label='Train Loss', alpha=0.8)
    ax.plot(history['val_loss'], label='Val Loss', alpha=0.8)
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Loss')
    ax.set_title('Training and Validation Loss')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # Plot 2: Accuracy
    ax = axes[0, 1]
    ax.plot(history['train_acc'], label='Train Acc', alpha=0.8)
    ax.plot(history['val_acc'], label='Val Acc', alpha=0.8)
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Accuracy (%)')
    ax.set_title('Training and Validation Accuracy')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # Plot 3: Confusion Matrix
    ax = axes[1, 0]
    cm = confusion_matrix(y_true, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=ax,
                xticklabels=class_names, yticklabels=class_names)
    ax.set_ylabel('True Label')
    ax.set_xlabel('Predicted Label')
    ax.set_title('Confusion Matrix')
    
    # Plot 4: ROC Curve
    ax = axes[1, 1]
    fpr, tpr, _ = roc_curve(y_true, y_probs)
    auc_score = roc_auc_score(y_true, y_probs)
    ax.plot(fpr, tpr, label=f'ROC Curve (AUC = {auc_score:.3f})', linewidth=2)
    ax.plot([0, 1], [0, 1], 'k--', label='Random Classifier')
    ax.set_xlabel('False Positive Rate')
    ax.set_ylabel('True Positive Rate')
    ax.set_title('ROC Curve')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('training_results.png', dpi=300, bbox_inches='tight')
    print("\nPlots saved as 'training_results.png'")
    
    return fig


def main():
    # Configuration
    CALLGRAPH_PATH = 'callgraph.txt'
    THRESHOLD = 2  # Functions with uses <= 2 are labeled as rarely used
    TEST_SIZE = 0.2
    VAL_SIZE = 0.2
    BATCH_SIZE = 32
    EPOCHS = 100
    LEARNING_RATE = 0.001
    RANDOM_STATE = 42
    
    # Set random seeds
    np.random.seed(RANDOM_STATE)
    torch.manual_seed(RANDOM_STATE)
    
    # Device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # Step 1: Parse callgraph
    parser = CallGraphParser(CALLGRAPH_PATH)
    parser.parse()
    
    # Step 2: Extract features
    df = parser.extract_features(threshold=THRESHOLD)
    
    # Step 3: Prepare data
    feature_columns = ['num_calls', 'in_degree', 'out_degree', 'pagerank',
                      'func_length', 'has_number', 'has_underscore', 'is_internal',
                      'has_sqlite_prefix']
    
    X = df[feature_columns].values
    y = df['label'].values
    
    # Normalize features
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    # Train/Val/Test split
    X_temp, X_test, y_temp, y_test = train_test_split(
        X_scaled, y, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=y
    )
    
    X_train, X_val, y_train, y_val = train_test_split(
        X_temp, y_temp, test_size=VAL_SIZE, random_state=RANDOM_STATE, stratify=y_temp
    )
    
    print(f"\nData split:")
    print(f"Training set: {len(X_train)} samples")
    print(f"Validation set: {len(X_val)} samples")
    print(f"Test set: {len(X_test)} samples")
    
    # Create datasets and dataloaders
    train_dataset = CodeDataset(X_train, y_train)
    val_dataset = CodeDataset(X_val, y_val)
    test_dataset = CodeDataset(X_test, y_test)
    
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
    
    # Step 4: Create model
    input_dim = X_train.shape[1]
    model = DeadCodeClassifier(input_dim=input_dim, hidden_dims=[64, 32, 16], dropout=0.3)
    
    print(f"\nModel architecture:")
    print(model)
    print(f"\nTotal parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    # Step 5: Train model
    trainer = Trainer(model, device=device)
    trainer.fit(train_loader, val_loader, epochs=EPOCHS, lr=LEARNING_RATE)
    
    # Step 6: Evaluate on test set
    print("\n" + "="*50)
    print("FINAL EVALUATION ON TEST SET")
    print("="*50)
    
    y_true, y_pred, y_probs = trainer.evaluate(test_loader)
    
    print("\nClassification Report:")
    print(classification_report(y_true, y_pred, 
                                target_names=['Frequently Used', 'Rarely Used'],
                                digits=4))
    
    print(f"\nROC-AUC Score: {roc_auc_score(y_true, y_probs):.4f}")
    
    # Step 7: Plot results
    plot_results(trainer.history, y_true, y_pred, y_probs)
    
    # Step 8: Save model and scaler
    torch.save({
        'model_state_dict': model.state_dict(),
        'scaler': scaler,
        'feature_columns': feature_columns,
        'threshold': THRESHOLD
    }, 'dead_code_detector.pth')
    
    print("\nModel saved as 'dead_code_detector.pth'")
    
    # Step 9: Example predictions
    print("\n" + "="*50)
    print("EXAMPLE PREDICTIONS")
    print("="*50)
    
    # Get some examples from each class
    rarely_used = df[df['label'] == 1].head(5)
    frequently_used = df[df['label'] == 0].head(5)
    
    for idx, row in rarely_used.iterrows():
        print(f"\n✗ {row['function_name']}")
        print(f"  Uses: {row['uses']}, Calls: {row['num_calls']}, Label: Rarely Used")
    
    for idx, row in frequently_used.iterrows():
        print(f"\n✓ {row['function_name']}")
        print(f"  Uses: {row['uses']}, Calls: {row['num_calls']}, Label: Frequently Used")
    
    print("\n" + "="*50)
    print("Training complete!")
    print("="*50)


if __name__ == "__main__":
    main()